diff --git a/.clang-format b/.clang-format new file mode 100644 index 00000000..2593ef54 --- /dev/null +++ b/.clang-format @@ -0,0 +1 @@ +BasedOnStyle: Google \ No newline at end of file diff --git a/.gitignore b/.gitignore index d6297828..667fd596 100644 --- a/.gitignore +++ b/.gitignore @@ -106,3 +106,21 @@ venv.bak/ # protobuf stuff tf/proto/ +**/*_pb2.py +**/*_pb2.pyi + +# Meson +builddir/ +subprojects/ + +# LLM Agents +# Use AGENTS.md; symlink other files to it if needed. +.claude/ +CLAUDE.local.md +.claude.json +CLAUDE.md +GEMINI.md + +# IDE +.vscode/ +.idea/ \ No newline at end of file diff --git a/.gitmodules b/.gitmodules index 850945c2..eb37789e 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,3 @@ -[submodule "libs/lczero-common"] - path = libs/lczero-common - url = https://github.com/LeelaChessZero/lczero-common.git \ No newline at end of file +[submodule "libs/lc0"] + path = libs/lc0 + url = https://github.com/LeelaChessZero/lc0.git diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..cf9dbbfa --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,59 @@ +# AGENTS.md + +This repository contains training script for the Leela Chess Zero project. +They are being rewritten. + +* Old code is located in the `tf/` directory. +* New python code is located in the `src/` directory. +* New C++ code is located in the `csrc/` directory. + +The old code is Python/TensorFlow-based, new code is Python/JAX-based with +modules written in C++, operating through pybind11. + +The build system for C++ code is meson. During development, the project is built +in the `builddir/`. + +## Testing and Building + +* C++ tests use GTest framework + * Do not insert Sleeps in tests, it slows down presubmit. Instead use e.g. + absl::Notification, or std::future +* Tests are defined in `meson.build` with `test()` function + * When debugging, don't forget to build them before running `meson test` or + `builddir/test` +* Run tests: `meson test -C builddir/` +* Python tests use `pytest` framework + * Do not add custom main function, exception catching to report errors, any + "test passed" messages etc. Use `pytest` fixtures and assertions. +* Build: `meson compile -C builddir/` from build directory +* Format code: `just format` +* There is a commit hook that runs `just pre-commit`, which runs tests and + checks formatting. You may want to run it before attempting to commit. +* We use Google C++ style guide. + * That means 80 columns. + * That means comments should be in full sentences with periods in the end. + * When conditional or loop fits one line, it must be written as one line + without braces, for example: + `if (condition) return value;` + * Prefer `absl` to `std` (e.g. `absl::c_` algorithms, `absl::Mutex`, + `absl::StrCat`, etc.) +* We use `uv` for Python package and venv management, and to running the + application. +* Run TUI app: `uv run tui --config=` + +* Do not attempt to run TUI — it messes up the Agent interface and session has + to be killed. Ask me to check it for you manually instead. + +* Do not commit unless explicitly asked. + +## IMPORTANT + +* NEVER add `# type: ignore` or other ways to mask/silence errors instead of + fixing them. +* Rely on protobuf default values. DO NOT write code like + `config.has_foo() ? config.foo() : default_value;` + +## Documentation + +* Documentation is in the `docs/` directory. +* The contents is in [The index](docs/index.md) diff --git a/README.md b/README.md index 367ddc1b..d818ebba 100644 --- a/README.md +++ b/README.md @@ -78,3 +78,16 @@ The training pipeline will automatically restore from a previous model if it exi ## Supervised training Generating trainingdata from pgn files is currently broken and has low priority, feel free to create a PR. + +## Building 2025-08 version. + +1. Make sure `uv` and `justfile` are installed (plus `meson` and other stuff, potentially `protoc`). +2. `git submodule update` +3. `uv venv` (!important! do this before running meson; otherwise meson will build module for wrong python) +4. `uv sync` +5. `CXX=clang++ CC=clang uv run meson setup build/release/ --buildtype=release --native-file=native.ini` (clang is optional, should build fine with default compiler) +6. `just build-proto` +7. `meson compile -C build/release/` +8. `ln -s -T ../../build/release/_lczero_training.cpython-311-x86_64-linux-gnu.so src/lczero_training/_lczero_training.so` + +14. Run it! `uv run lc0-tui --config docs/example.textproto` diff --git a/csrc/loader/chunk_source/chunk_source.h b/csrc/loader/chunk_source/chunk_source.h new file mode 100644 index 00000000..cf8fda46 --- /dev/null +++ b/csrc/loader/chunk_source/chunk_source.h @@ -0,0 +1,35 @@ +#pragma once + +#include +#include +#include +#include + +#include "loader/frame_type.h" + +namespace lczero { +namespace training { + +// Interface for providing training data chunks. +// A chunk source provides access to one or more chunks of training data. +// It's assumed that all chunks in a source for one group for sorting purposes, +// therefore GetChunkSortKey() returns just one key for the entire source. This +// allows to know the key before reading/indexing the chunks. +class ChunkSource { + public: + virtual ~ChunkSource() = default; + + // Returns a sort key (e.g. filename or a timestamp). + virtual std::string GetChunkSortKey() const = 0; + + // Returns the number of chunks in this source. + virtual size_t GetChunkCount() const = 0; + + // Returns the data for the chunk at the given index. Returns std::nullopt if + // the chunk could not be read or if the data size is not a multiple of the + // expected frame size. + virtual std::optional> GetChunkData(size_t index) = 0; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/chunk_source/chunk_source_view.h b/csrc/loader/chunk_source/chunk_source_view.h new file mode 100644 index 00000000..90a7aee4 --- /dev/null +++ b/csrc/loader/chunk_source/chunk_source_view.h @@ -0,0 +1,47 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "loader/chunk_source/chunk_source.h" + +namespace lczero { +namespace training { + +// ChunkSourceView provides a view over another ChunkSource. +// It exposes a remapped subset/order of chunks defined by indices into the +// underlying source. It does not own or copy the data; it forwards calls to +// the wrapped source. +class ChunkSourceView : public ChunkSource { + public: + // Constructs a view over an existing chunk source. The indices vector maps + // local indices in the view to indices of the underlying source. + ChunkSourceView(std::shared_ptr source, + std::vector indices) + : source_(std::move(source)), indices_(std::move(indices)) {} + + ~ChunkSourceView() override = default; + + private: + std::string GetChunkSortKey() const override { + return source_->GetChunkSortKey(); + } + + size_t GetChunkCount() const override { return indices_.size(); } + + std::optional> GetChunkData(size_t index) override { + if (index >= indices_.size()) return std::nullopt; + const size_t src_index = static_cast(indices_[index]); + return source_->GetChunkData(src_index); + } + + std::shared_ptr source_; + std::vector indices_; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/chunk_source/debug_chunk_source.cc b/csrc/loader/chunk_source/debug_chunk_source.cc new file mode 100644 index 00000000..f2058e35 --- /dev/null +++ b/csrc/loader/chunk_source/debug_chunk_source.cc @@ -0,0 +1,57 @@ +#include "loader/chunk_source/debug_chunk_source.h" + +#include +#include +#include +#include +#include +#include + +#include "absl/hash/hash.h" +#include "absl/strings/str_format.h" + +namespace lczero { +namespace training { + +DebugChunkSource::DebugChunkSource(uint64_t id, double mean_chunk_count) + : id_(id), mean_chunk_count_(mean_chunk_count) {} + +std::string DebugChunkSource::GetChunkSortKey() const { + return absl::StrFormat("%08" PRIu64, id_); +} + +size_t DebugChunkSource::GetChunkCount() const { + if (!cached_chunk_count_.has_value()) { + std::mt19937_64 rng(id_); + const double stddev = std::max(1.0, mean_chunk_count_ / 4.0); + std::normal_distribution distribution(mean_chunk_count_, stddev); + const double sampled = distribution(rng); + const auto rounded = + static_cast(std::llround(std::max(sampled, 1.0))); + cached_chunk_count_ = static_cast(rounded); + } + return *cached_chunk_count_; +} + +std::optional> DebugChunkSource::GetChunkData( + size_t index) { + const auto seed_pair = std::make_pair(id_, index); + const uint64_t seed = static_cast( + absl::Hash>{}(seed_pair)); + std::mt19937_64 rng(seed); + std::uniform_int_distribution frame_count_distribution(1, 200); + const int frame_count = frame_count_distribution(rng); + + std::vector result(frame_count); + for (int frame_index = 0; frame_index < frame_count; ++frame_index) { + result[frame_index] = FrameType{}; + result[frame_index].planes[0] = static_cast(id_); + result[frame_index].planes[1] = static_cast(index); + result[frame_index].planes[2] = static_cast(frame_index); + } + + return result; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/chunk_source/debug_chunk_source.h b/csrc/loader/chunk_source/debug_chunk_source.h new file mode 100644 index 00000000..b93ea0a7 --- /dev/null +++ b/csrc/loader/chunk_source/debug_chunk_source.h @@ -0,0 +1,41 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "loader/chunk_source/chunk_source.h" + +namespace lczero { +namespace training { + +// DebugChunkSource synthesizes deterministic pseudo-random chunks for loader +// debugging. Each instance is identified by an integer id. The class produces +// a chunk count sampled from a normal distribution with the provided mean and +// mean / 4 standard deviation. The id serves as the seed, which keeps the +// number of chunks stable across runs. Individual chunks contain a +// pseudo-random number of FrameType frames (between one and 200) that are +// generated on demand. The generation seed depends on both the source id and +// chunk index. This lets shuffling logic exercise variable chunk sizes while +// keeping the content reproducible. Each generated frame is zero-initialized, +// but the first three entries of the planes array encode, respectively, the +// source id, the chunk index, and the frame index within the chunk. This makes +// it easy to reason about ordering and grouping when inspecting chunk payloads. +class DebugChunkSource : public ChunkSource { + public: + DebugChunkSource(uint64_t id, double mean_chunk_count); + + private: + std::string GetChunkSortKey() const override; + size_t GetChunkCount() const override; + std::optional> GetChunkData(size_t index) override; + + uint64_t id_; + double mean_chunk_count_; + mutable std::optional cached_chunk_count_; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/chunk_source/rawfile_chunk_source.cc b/csrc/loader/chunk_source/rawfile_chunk_source.cc new file mode 100644 index 00000000..77fd2879 --- /dev/null +++ b/csrc/loader/chunk_source/rawfile_chunk_source.cc @@ -0,0 +1,59 @@ +#include "loader/chunk_source/rawfile_chunk_source.h" + +#include + +#include +#include + +#include "trainingdata/trainingdata_v6.h" +#include "utils/files.h" +#include "utils/gz.h" + +namespace lczero { +namespace training { + +RawFileChunkSource::RawFileChunkSource( + const std::filesystem::path& filename, + ChunkSourceLoaderConfig::FrameFormat frame_format) + : filename_(filename), frame_format_(frame_format) {} + +RawFileChunkSource::~RawFileChunkSource() = default; + +std::string RawFileChunkSource::GetChunkSortKey() const { + return std::filesystem::path(filename_).filename().string(); +} + +size_t RawFileChunkSource::GetChunkCount() const { return 1; } + +std::optional> RawFileChunkSource::GetChunkData( + size_t index) { + if (index != 0) return std::nullopt; + std::string data = ReadFileToString(filename_); + if (data.empty()) return std::nullopt; + + const size_t input_size = + frame_format_ == ChunkSourceLoaderConfig::V7TrainingData + ? sizeof(V7TrainingData) + : sizeof(V6TrainingData); + if (data.size() % input_size != 0) { + LOG(WARNING) << "File " << filename_ << " size " << data.size() + << " is not a multiple of input frame size " << input_size; + return std::nullopt; + } + + const size_t num_frames = data.size() / input_size; + std::vector result(num_frames); + + if (frame_format_ == ChunkSourceLoaderConfig::V7TrainingData) { + std::memcpy(result.data(), data.data(), data.size()); + } else { + const auto* v6_data = reinterpret_cast(data.data()); + for (size_t i = 0; i < num_frames; ++i) { + std::memcpy(&result[i], &v6_data[i], sizeof(V6TrainingData)); + } + } + return result; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/chunk_source/rawfile_chunk_source.h b/csrc/loader/chunk_source/rawfile_chunk_source.h new file mode 100644 index 00000000..426c1143 --- /dev/null +++ b/csrc/loader/chunk_source/rawfile_chunk_source.h @@ -0,0 +1,30 @@ +#pragma once + +#include +#include + +#include "loader/chunk_source/chunk_source.h" +#include "proto/data_loader_config.pb.h" + +namespace lczero { +namespace training { + +// A chunk source that reads a single (potentially gzipped) file as a single +// chunk. +class RawFileChunkSource : public ChunkSource { + public: + RawFileChunkSource(const std::filesystem::path& filename, + ChunkSourceLoaderConfig::FrameFormat frame_format); + ~RawFileChunkSource(); + + private: + std::string GetChunkSortKey() const override; + size_t GetChunkCount() const override; + std::optional> GetChunkData(size_t index) override; + + std::string filename_; + ChunkSourceLoaderConfig::FrameFormat frame_format_; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/chunk_source/tar_chunk_source.cc b/csrc/loader/chunk_source/tar_chunk_source.cc new file mode 100644 index 00000000..d87d6e99 --- /dev/null +++ b/csrc/loader/chunk_source/tar_chunk_source.cc @@ -0,0 +1,256 @@ +#include "loader/chunk_source/tar_chunk_source.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "trainingdata/trainingdata_v6.h" +#include "utils/gz.h" + +namespace lczero { +namespace training { +namespace { +struct TarHeader { + std::array name; + std::array mode; + std::array uid; + std::array gid; + std::array size; + std::array mtime; + std::array chksum; + uint8_t typeflag; + std::array linkname; + std::array magic; + std::array version; + std::array uname; + std::array gname; + std::array devmajor; + std::array devminor; + std::array prefix; + std::array padding; +}; +static_assert(sizeof(TarHeader) == 512, "TarHeader must be exactly 512 bytes"); + +uint64_t ParseOctal(const std::array& octal) { + uint64_t value = 0; + for (uint8_t digit : octal) { + if (!digit) break; + value = (value << 3) + (digit - '0'); + } + return value; +} + +bool ReadExact(int fd, off_t offset, void* buffer, size_t size) { + char* out = static_cast(buffer); + size_t read_total = 0; + while (read_total < size) { + const ssize_t read_now = + pread(fd, out + read_total, size - read_total, offset + read_total); + if (read_now <= 0) return false; + read_total += static_cast(read_now); + } + return true; +} + +std::optional ReadGzipPrefix(int fd, off_t offset, size_t size, + size_t max_bytes) { + if (max_bytes == 0) return std::string(); + + z_stream strm = {}; + if (inflateInit2(&strm, 16 + MAX_WBITS) != Z_OK) { + return std::nullopt; + } + + constexpr size_t kChunkSize = 16384; + std::array input_buffer; + std::array output_buffer; + + std::string output; + output.reserve(std::min(max_bytes, kChunkSize)); + + size_t remaining = size; + off_t current_offset = offset; + bool finished = false; + + while (remaining > 0 && !finished && output.size() < max_bytes) { + const size_t to_read = std::min(remaining, kChunkSize); + if (!ReadExact(fd, current_offset, input_buffer.data(), to_read)) { + inflateEnd(&strm); + return std::nullopt; + } + remaining -= to_read; + current_offset += static_cast(to_read); + + strm.next_in = reinterpret_cast(input_buffer.data()); + strm.avail_in = static_cast(to_read); + + while (strm.avail_in > 0 && output.size() < max_bytes) { + strm.next_out = reinterpret_cast(output_buffer.data()); + strm.avail_out = kChunkSize; + + const int ret = inflate(&strm, Z_NO_FLUSH); + if (ret == Z_STREAM_ERROR || ret == Z_NEED_DICT || ret == Z_DATA_ERROR || + ret == Z_MEM_ERROR) { + inflateEnd(&strm); + return std::nullopt; + } + + const size_t produced = kChunkSize - strm.avail_out; + const size_t to_copy = std::min(produced, max_bytes - output.size()); + output.append(output_buffer.data(), to_copy); + + if (ret == Z_STREAM_END) { + finished = true; + break; + } + } + } + + inflateEnd(&strm); + return output; +} + +} // namespace + +TarChunkSource::TarChunkSource( + const std::filesystem::path& filename, + ChunkSourceLoaderConfig::FrameFormat frame_format) + : path_(filename), + filename_(filename.filename().string()), + frame_format_(frame_format) { + fd_ = open(path_.c_str(), O_RDONLY | O_CLOEXEC); + if (fd_ < 0) { + throw std::runtime_error( + absl::StrCat("Failed to open tar file: ", path_.string(), ": ", + std::strerror(errno))); + } + // Perform indexing during construction. + Index(); +} + +TarChunkSource::~TarChunkSource() { Close(); } + +void TarChunkSource::Close() { + if (fd_ >= 0 && close(fd_) != 0) { + PLOG(WARNING) << "Failed to close tar file descriptor for " << path_; + } + fd_ = -1; +} + +std::string TarChunkSource::GetChunkSortKey() const { return filename_; } + +void TarChunkSource::Index() { + assert(files_.empty()); + + off_t offset = 0; + + while (true) { + TarHeader header; + if (!ReadExact(fd_, offset, &header, sizeof(header))) { + LOG(WARNING) << "Truncated tar file: " << filename_; + break; + } + offset += sizeof(header); + + if (header.name[0] == '\0') break; // End of file. + + switch (header.typeflag) { + case '5': // Directory + continue; + case '0': // Regular file + break; + default: + LOG(WARNING) << "Unsupported tar header type: " << header.typeflag; + continue; + } + + std::string_view fname(const_cast(header.name.data())); + const std::filesystem::path filepath = std::filesystem::path(fname); + const off_t file_offset = offset; + const size_t size = ParseOctal(header.size); + const size_t padded_size = ((size + 511) / 512) * 512; + offset = file_offset + static_cast(padded_size); + + if (filepath.filename() == "LICENSE") continue; + files_.push_back({file_offset, size, filepath.extension() == ".gz"}); + } + + LOG(INFO) << "Read " << files_.size() << " entries from " << filename_; +} + +size_t TarChunkSource::GetChunkCount() const { return files_.size(); } + +std::optional> TarChunkSource::GetChunkData( + size_t index) { + if (index >= files_.size()) { + throw std::out_of_range("File index out of range"); + } + const auto& file_entry = files_[index]; + std::string content(file_entry.size, '\0'); + if (!ReadExact(fd_, file_entry.offset, content.data(), file_entry.size)) { + return std::nullopt; + } + if (file_entry.is_gzip) { + try { + content = GunzipBuffer(content); + } catch (const GunzipError& e) { + return std::nullopt; + } + } + if (content.empty()) return std::nullopt; + + const size_t input_size = + frame_format_ == ChunkSourceLoaderConfig::V7TrainingData + ? sizeof(V7TrainingData) + : sizeof(V6TrainingData); + if (content.size() % input_size != 0) { + LOG(WARNING) << "Chunk " << index << " from " << filename_ << " size " + << content.size() << " is not a multiple of input frame size " + << input_size; + return std::nullopt; + } + + const size_t num_frames = content.size() / input_size; + std::vector result(num_frames); + + if (frame_format_ == ChunkSourceLoaderConfig::V7TrainingData) { + std::memcpy(result.data(), content.data(), content.size()); + } else { + const auto* v6_data = + reinterpret_cast(content.data()); + for (size_t i = 0; i < num_frames; ++i) { + std::memcpy(&result[i], &v6_data[i], sizeof(V6TrainingData)); + } + } + return result; +} + +std::optional TarChunkSource::GetChunkPrefix(size_t index, + size_t max_bytes) { + if (index >= files_.size()) { + throw std::out_of_range("File index out of range"); + } + const auto& file_entry = files_[index]; + if (file_entry.is_gzip) { + return ReadGzipPrefix(fd_, file_entry.offset, file_entry.size, max_bytes); + } + + const size_t to_read = std::min(file_entry.size, max_bytes); + std::string content(to_read, '\0'); + if (!ReadExact(fd_, file_entry.offset, content.data(), to_read)) { + return std::nullopt; + } + return content; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/chunk_source/tar_chunk_source.h b/csrc/loader/chunk_source/tar_chunk_source.h new file mode 100644 index 00000000..364db864 --- /dev/null +++ b/csrc/loader/chunk_source/tar_chunk_source.h @@ -0,0 +1,46 @@ +#pragma once + +#include + +#include +#include +#include + +#include "loader/chunk_source/chunk_source.h" +#include "proto/data_loader_config.pb.h" + +namespace lczero { +namespace training { + +// A chunk source that reads a tar archive and provides access to its files as +// chunks. Each file in the tar is treated as a separate chunk. +class TarChunkSource : public ChunkSource { + public: + TarChunkSource(const std::filesystem::path& filename, + ChunkSourceLoaderConfig::FrameFormat frame_format); + ~TarChunkSource() override; + std::string GetChunkSortKey() const override; + size_t GetChunkCount() const override; + std::optional> GetChunkData(size_t index) override; + std::optional GetChunkPrefix(size_t index, size_t max_bytes); + + private: + // Performs one-time indexing during construction. Not part of the interface. + void Index(); + struct FileEntry { + off_t offset; + size_t size; + bool is_gzip; + }; + + void Close(); + + int fd_ = -1; + std::vector files_; + std::filesystem::path path_; + std::string filename_; + ChunkSourceLoaderConfig::FrameFormat frame_format_; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/data_loader.cc b/csrc/loader/data_loader.cc new file mode 100644 index 00000000..9bef3beb --- /dev/null +++ b/csrc/loader/data_loader.cc @@ -0,0 +1,242 @@ +#include "loader/data_loader.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "loader/data_loader_metrics.h" + +namespace lczero { +namespace training { +DataLoaderConfig DataLoader::ParseConfig(const std::string& serialized_config) { + DataLoaderConfig config; + config.ParseFromString(serialized_config); + return config; +} + +DataLoader::DataLoader(const std::string& serialized_data_loader_config) + : metrics_aggregator_( + [](DataLoaderMetricsProto& m) { m.Clear(); }, + [](DataLoaderMetricsProto& dest, const DataLoaderMetricsProto& src) { + UpdateFrom(dest, src); + }) { + DataLoaderConfig config = ParseConfig(serialized_data_loader_config); + AddStages(config); + BuildOutputMapping(config); + LOG(INFO) << "DataLoader initialized with " << stage_registry_.size() + << " stage(s) and " << outputs_.size() << " output(s)."; +} + +DataLoader::~DataLoader() { Stop(); } + +void DataLoader::AddStages(const std::string& serialized_data_loader_config) { + AddStages(ParseConfig(serialized_data_loader_config)); +} + +void DataLoader::AddStages(const DataLoaderConfig& config) { + for (const auto& stage_config : config.stage()) AddStage(stage_config); + for (const auto& stage_config : config.stage()) SetStageInputs(stage_config); +} + +void DataLoader::AddStage(const StageConfig& stage_config) { + if (started_) { + throw std::runtime_error("Cannot add stages after DataLoader has started."); + } + + auto stage = CreateStage(stage_config); + if (!stage_config.has_name()) { + throw std::runtime_error("Stage configuration is missing name."); + } + + LOG(INFO) << "Adding stage '" << stage_config.name() << "'."; + stage_registry_.AddStage(stage_config.name(), std::move(stage)); +} + +void DataLoader::SetStageInputs(const StageConfig& stage_config) { + // Resolve input names to queue pointers. + std::vector input_queues; + input_queues.reserve(stage_config.input_size()); + for (const auto& input_name : stage_config.input()) { + QueueBase* queue = stage_registry_.GetStageOutput(input_name); + if (!queue) { + throw std::runtime_error(absl::StrCat("Input stage '", input_name, + "' not found for stage '", + stage_config.name(), "'.")); + } + input_queues.push_back(queue); + } + + // Wire up inputs. + auto it = absl::c_find_if(stage_registry_.stages(), [&](const auto& p) { + return p.first == stage_config.name(); + }); + if (it == stage_registry_.stages().end()) { + throw std::runtime_error(absl::StrCat("Stage '", stage_config.name(), + "' not found in registry.")); + } + it->second->SetInputs(absl::MakeSpan(input_queues)); +} + +void DataLoader::Start() { + if (started_) { + throw std::runtime_error("DataLoader has already been started."); + } + for (auto& [name, stage] : stage_registry_.stages()) { + LOG(INFO) << "Starting stage '" << name << "'."; + stage->Start(); + } + + metrics_thread_ = std::jthread( + [this](std::stop_token stop_token) { MetricsThread(stop_token); }); + started_ = true; + stopped_ = false; + LOG(INFO) << "DataLoader started."; +} + +TensorTuple DataLoader::GetNext(std::string_view alias) { + Queue* q = GetOutputQueue(alias); + if (!q) { + std::string alias_list = absl::StrJoin( + outputs_, ", ", + [](std::string* out, const auto& p) { absl::StrAppend(out, p.first); }); + throw std::runtime_error(absl::StrCat("Unknown DataLoader output: '", alias, + "'. Available outputs: [", alias_list, + "].")); + } + return q->Get(); +} + +std::optional DataLoader::MaybeGetNext(std::string_view alias) { + Queue* q = GetOutputQueue(alias); + if (!q) { + std::string alias_list = absl::StrJoin( + outputs_, ", ", + [](std::string* out, const auto& p) { absl::StrAppend(out, p.first); }); + throw std::runtime_error(absl::StrCat("Unknown DataLoader output: '", alias, + "'. Available outputs: [", alias_list, + "].")); + } + return q->MaybeGet(); +} + +void DataLoader::Stop() { + if (stopped_) return; + + LOG(INFO) << "Stopping DataLoader."; + + if (metrics_thread_.joinable()) { + metrics_thread_.request_stop(); + metrics_thread_.join(); + } + + for (auto& [name, stage] : stage_registry_.stages()) { + LOG(INFO) << "Stopping stage '" << name << "'."; + stage->Stop(); + } + + stopped_ = true; + started_ = false; + LOG(INFO) << "DataLoader stopped."; +} + +std::pair DataLoader::GetBucketMetrics( + int time_period, bool include_pending) const { + auto [metrics, duration] = metrics_aggregator_.GetBucketMetrics( + static_cast(time_period), + include_pending ? std::make_optional(std::chrono::steady_clock::now()) + : std::nullopt); + float duration_seconds = std::chrono::duration(duration).count(); + return {metrics.OutputAsString(), duration_seconds}; +} + +std::pair DataLoader::GetAggregateEndingNow( + float duration_seconds, bool include_pending) const { + std::chrono::nanoseconds duration_ns = + std::isinf(duration_seconds) + ? std::chrono::nanoseconds::max() + : std::chrono::nanoseconds( + static_cast(duration_seconds * 1e9)); + auto [metrics, duration] = metrics_aggregator_.GetAggregateEndingNow( + duration_ns, include_pending + ? std::make_optional(std::chrono::steady_clock::now()) + : std::nullopt); + float result_duration_seconds = + std::chrono::duration(duration).count(); + return {metrics.OutputAsString(), result_duration_seconds}; +} + +void DataLoader::MetricsThread(std::stop_token stop_token) { + while (!stop_token.stop_requested()) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + DataLoaderMetricsProto metrics; + for (auto& [name, stage] : stage_registry_.stages()) { + StageMetricProto stage_metric = stage->FlushMetrics(); + stage_metric.set_name(name); + *metrics.add_stage_metrics() = std::move(stage_metric); + } + + metrics_aggregator_.RecordMetrics(std::move(metrics)); + metrics_aggregator_.Advance(std::chrono::steady_clock::now()); + } + LOG(INFO) << "Metrics thread stopping."; +} + +std::vector> +DataLoader::SendControlMessage(const StageControlRequest& request) { + std::vector> responses; + responses.reserve(stage_registry_.size()); + for (auto& [name, stage] : stage_registry_.stages()) { + std::optional response = stage->Control(request); + if (response.has_value()) { + responses.emplace_back(name, std::move(*response)); + } + } + return responses; +} + +void DataLoader::BuildOutputMapping(const DataLoaderConfig& config) { + outputs_.clear(); + outputs_.reserve(config.output_size()); + + for (const auto& out_spec : config.output()) { + auto it = absl::c_find(out_spec, ':'); + std::string alias = it == out_spec.end() + ? std::string("") + : std::string(out_spec.begin(), it); + std::string stage_name = it == out_spec.end() + ? std::string(out_spec) + : std::string(it + 1, out_spec.end()); + + // Ensure alias is unique. + if (absl::c_find_if(outputs_, [&](const auto& p) { + return p.first == alias; + }) != outputs_.end()) { + throw std::runtime_error( + absl::StrCat("Duplicate output alias specified: ", alias)); + } + + Queue* queue = + stage_registry_.GetTypedStageOutput(stage_name); + if (queue == nullptr) { + throw std::runtime_error( + absl::StrCat("Output stage not found or wrong type: ", stage_name)); + } + outputs_.emplace_back(std::move(alias), queue); + } +} + +Queue* DataLoader::GetOutputQueue(std::string_view alias) const { + auto it = absl::c_find_if(outputs_, + [&](const auto& p) { return p.first == alias; }); + if (it == outputs_.end()) return nullptr; + return it->second; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/data_loader.h b/csrc/loader/data_loader.h new file mode 100644 index 00000000..71f9661b --- /dev/null +++ b/csrc/loader/data_loader.h @@ -0,0 +1,65 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "loader/stages/stage_factory.h" +#include "proto/data_loader_config.pb.h" +#include "proto/stage_control.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/metrics/exponential_aggregator.h" +#include "utils/queue.h" +#include "utils/tensor.h" + +namespace lczero { +namespace training { + +using TensorDict = + std::vector>>; + +class DataLoader { + public: + using MetricsAggregator = ExponentialAggregator; + + explicit DataLoader(const std::string& serialized_data_loader_config); + ~DataLoader(); + + void Start(); + TensorTuple GetNext(std::string_view alias); + std::optional MaybeGetNext(std::string_view alias); + void Stop(); + std::pair GetBucketMetrics(int time_period, + bool include_pending) const; + std::pair GetAggregateEndingNow( + float duration_seconds, bool include_pending) const; + + void AddStages(const DataLoaderConfig& config); + void AddStages(const std::string& serialized_data_loader_config); + + std::vector> SendControlMessage( + const StageControlRequest& request); + + private: + void AddStage(const StageConfig& stage_config); + void SetStageInputs(const StageConfig& stage_config); + static DataLoaderConfig ParseConfig( + const std::string& serialized_data_loader_config); + void MetricsThread(std::stop_token stop_token); + void BuildOutputMapping(const DataLoaderConfig& config); + Queue* GetOutputQueue(std::string_view alias) const; + + StageRegistry stage_registry_; + std::vector*>> outputs_; + MetricsAggregator metrics_aggregator_; + std::jthread metrics_thread_; + bool started_ = false; + bool stopped_ = false; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/data_loader_metrics.cc b/csrc/loader/data_loader_metrics.cc new file mode 100644 index 00000000..551bad76 --- /dev/null +++ b/csrc/loader/data_loader_metrics.cc @@ -0,0 +1,124 @@ +// ABOUTME: Implementation of UpdateFrom functions for data loader metric +// protobuf messages. ABOUTME: Handles aggregation of generic stage metrics. + +#include "loader/data_loader_metrics.h" + +#include + +#include "absl/strings/string_view.h" +#include "utils/metrics/statistics_metric.h" + +namespace lczero { +namespace training { +namespace { + +template +ProtoT* FindByName(std::vector* entries, absl::string_view name) { + for (auto& entry : *entries) { + if (entry.name() == name) return &entry; + } + return nullptr; +} + +} // namespace + +void UpdateFrom(QueueMetricProto& dest, const QueueMetricProto& src) { + if (src.has_name()) dest.set_name(src.name()); + dest.set_put_count(dest.put_count() + src.put_count()); + dest.set_get_count(dest.get_count() + src.get_count()); + dest.set_drop_count(dest.drop_count() + src.drop_count()); + UpdateFrom(*dest.mutable_queue_fullness(), src.queue_fullness()); + if (src.has_queue_capacity()) dest.set_queue_capacity(src.queue_capacity()); +} + +void UpdateFrom(CountMetricProto& dest, const CountMetricProto& src) { + if (src.has_name()) dest.set_name(src.name()); + if (src.has_count()) dest.set_count(dest.count() + src.count()); +} + +void UpdateFrom(GaugeMetricProto& dest, const GaugeMetricProto& src) { + if (src.has_name()) dest.set_name(src.name()); + if (src.has_value()) dest.set_value(src.value()); + if (src.has_capacity()) dest.set_capacity(src.capacity()); +} + +void UpdateFrom(StageMetricProto& dest, const StageMetricProto& src) { + if (src.has_name()) dest.set_name(src.name()); + + for (const auto& load_metrics : src.load_metrics()) { + LoadMetricProto* dest_load = + load_metrics.has_name() + ? FindByName(dest.mutable_load_metrics(), load_metrics.name()) + : nullptr; + if (dest_load == nullptr) { + dest_load = dest.add_load_metrics(); + } + UpdateFrom(*dest_load, load_metrics); + } + + for (const auto& queue_metrics : src.queue_metrics()) { + QueueMetricProto* dest_queue = + queue_metrics.has_name() + ? FindByName(dest.mutable_queue_metrics(), queue_metrics.name()) + : nullptr; + if (dest_queue == nullptr) { + dest_queue = dest.add_queue_metrics(); + } + UpdateFrom(*dest_queue, queue_metrics); + } + + for (const auto& count_metrics : src.count_metrics()) { + CountMetricProto* dest_count = + count_metrics.has_name() + ? FindByName(dest.mutable_count_metrics(), count_metrics.name()) + : nullptr; + if (dest_count == nullptr) { + dest_count = dest.add_count_metrics(); + } + UpdateFrom(*dest_count, count_metrics); + } + + for (const auto& gauge_metrics : src.gauge_metrics()) { + GaugeMetricProto* dest_gauge = + gauge_metrics.has_name() + ? FindByName(dest.mutable_gauge_metrics(), gauge_metrics.name()) + : nullptr; + if (dest_gauge == nullptr) { + dest_gauge = dest.add_gauge_metrics(); + } + UpdateFrom(*dest_gauge, gauge_metrics); + } + + for (const auto& statistics_metrics : src.statistics_metrics()) { + StatisticsProtoDouble* dest_stats = + statistics_metrics.has_name() + ? FindByName(dest.mutable_statistics_metrics(), + statistics_metrics.name()) + : nullptr; + if (dest_stats == nullptr) { + dest_stats = dest.add_statistics_metrics(); + } + UpdateFrom(*dest_stats, statistics_metrics); + } + + if (src.has_last_chunk_key()) dest.set_last_chunk_key(src.last_chunk_key()); + if (src.has_anchor()) dest.set_anchor(src.anchor()); +} + +void UpdateFrom(DataLoaderMetricsProto& dest, + const DataLoaderMetricsProto& src) { + for (const auto& stage_metrics : src.stage_metrics()) { + StageMetricProto* dest_stage = + stage_metrics.has_name() + ? FindByName(dest.mutable_stage_metrics(), stage_metrics.name()) + : nullptr; + if (dest_stage == nullptr) { + dest_stage = dest.add_stage_metrics(); + } + + UpdateFrom(*dest_stage, stage_metrics); + } +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/data_loader_metrics.h b/csrc/loader/data_loader_metrics.h new file mode 100644 index 00000000..84e4ebb1 --- /dev/null +++ b/csrc/loader/data_loader_metrics.h @@ -0,0 +1,36 @@ +// ABOUTME: Header for UpdateFrom functions for data loader metric protobuf +// messages. ABOUTME: Declares functions for aggregating FilePathProvider and +// DataLoader metrics. + +#pragma once + +#include "absl/strings/string_view.h" +#include "proto/training_metrics.pb.h" +#include "utils/metrics/load_metric.h" +#include "utils/metrics/statistics_metric.h" +#include "utils/queue.h" + +namespace lczero { +namespace training { + +void UpdateFrom(QueueMetricProto& dest, const QueueMetricProto& src); +void UpdateFrom(CountMetricProto& dest, const CountMetricProto& src); +void UpdateFrom(GaugeMetricProto& dest, const GaugeMetricProto& src); +void UpdateFrom(StageMetricProto& dest, const StageMetricProto& src); +void UpdateFrom(DataLoaderMetricsProto& dest, + const DataLoaderMetricsProto& src); + +template +QueueMetricProto MetricsFromQueue(absl::string_view name, Queue& queue) { + QueueMetricProto result; + result.set_name(std::string(name)); + result.set_put_count(queue.GetTotalPutCount(true)); + result.set_get_count(queue.GetTotalGetCount(true)); + result.set_drop_count(queue.GetTotalDropCount(true)); + AddSample(*result.mutable_queue_fullness(), queue.Size()); + result.set_queue_capacity(queue.Capacity()); + return result; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/data_loader_test.cc b/csrc/loader/data_loader_test.cc new file mode 100644 index 00000000..f5ab9597 --- /dev/null +++ b/csrc/loader/data_loader_test.cc @@ -0,0 +1,33 @@ +#include "loader/data_loader.h" + +#include + +#include + +namespace lczero { +namespace training { + +TEST(DataLoaderTest, AllowsNoOutputsConfigured) { + DataLoaderConfig config; + auto* file_stage = config.add_stage(); + file_stage->set_name("file_path_provider"); + file_stage->mutable_file_path_provider()->set_directory("."); + + EXPECT_NO_THROW(DataLoader(config.OutputAsString())); +} + +TEST(DataLoaderTest, ThrowsOnDuplicateStageName) { + DataLoaderConfig config; + auto* first_stage = config.add_stage(); + first_stage->set_name("duplicate"); + first_stage->mutable_file_path_provider()->set_directory("."); + + auto* second_stage = config.add_stage(); + second_stage->set_name("duplicate"); + second_stage->mutable_file_path_provider()->set_directory("."); + + EXPECT_THROW(DataLoader(config.OutputAsString()), std::runtime_error); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/frame_type.h b/csrc/loader/frame_type.h new file mode 100644 index 00000000..25225ebf --- /dev/null +++ b/csrc/loader/frame_type.h @@ -0,0 +1,38 @@ +/* + This file is part of Leela Chess Zero. + Copyright (C) 2025 The LCZero Authors + + Leela Chess is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + Leela Chess is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with Leela Chess. If not, see . + + Additional permission under GNU GPL version 3 section 7 + + If you modify this Program, or any covered work, by linking or + combining it with NVIDIA Corporation's libraries from the NVIDIA CUDA + Toolkit and the NVIDIA CUDA Deep Neural Network library (or a + modified version of those libraries), containing parts covered by the + terms of the respective license agreement, the licensors of this + Program grant you additional permission to convey the resulting work. +*/ + +#pragma once + +#include "trainingdata/trainingdata_v7.h" + +namespace lczero { +namespace training { + +using FrameType = V7TrainingData; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/loader_main.cpp b/csrc/loader/loader_main.cpp new file mode 100644 index 00000000..53daf905 --- /dev/null +++ b/csrc/loader/loader_main.cpp @@ -0,0 +1,121 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "data_loader.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/metrics/printer.h" + +ABSL_FLAG(std::string, directory, "/home/crem/tmp/2025-07/lczero-training/", + "Directory to watch for training data files"); +ABSL_FLAG(size_t, chunk_pool_size, 1000000, "Size of the chunk shuffle buffer"); +ABSL_FLAG(size_t, reservoir_size_per_thread, 1000000, + "Size of the reservoir for frame sampling per thread"); + +namespace lczero { +namespace training { + +void Run() { + DataLoaderConfig config; + + // Configure file path provider stage. + auto* file_stage = config.add_stage(); + file_stage->set_name("file_path_provider"); + auto* file_path_provider = file_stage->mutable_file_path_provider(); + file_path_provider->set_directory(absl::GetFlag(FLAGS_directory)); + + // Configure chunk source loader stage. + auto* chunk_loader_stage = config.add_stage(); + chunk_loader_stage->set_name("chunk_source_loader"); + chunk_loader_stage->add_input(file_stage->name()); + chunk_loader_stage->mutable_chunk_source_loader(); + + // Configure shuffling chunk pool stage. + auto* chunk_pool_stage = config.add_stage(); + chunk_pool_stage->set_name("shuffling_chunk_pool"); + chunk_pool_stage->add_input(chunk_loader_stage->name()); + auto* shuffling_chunk_pool = chunk_pool_stage->mutable_shuffling_chunk_pool(); + shuffling_chunk_pool->set_chunk_pool_size( + absl::GetFlag(FLAGS_chunk_pool_size)); + + // Configure chunk unpacker stage. + auto* unpacker_stage = config.add_stage(); + unpacker_stage->set_name("chunk_unpacker"); + unpacker_stage->add_input(chunk_pool_stage->name()); + unpacker_stage->mutable_chunk_unpacker(); + + // Configure shuffling frame sampler stage. + auto* sampler_stage = config.add_stage(); + sampler_stage->set_name("shuffling_frame_sampler"); + sampler_stage->add_input(unpacker_stage->name()); + auto* shuffling_frame_sampler = + sampler_stage->mutable_shuffling_frame_sampler(); + shuffling_frame_sampler->set_reservoir_size_per_thread( + absl::GetFlag(FLAGS_reservoir_size_per_thread)); + + // Configure tensor generator stage. + auto* tensor_stage = config.add_stage(); + tensor_stage->set_name("tensor_generator"); + tensor_stage->add_input(sampler_stage->name()); + tensor_stage->mutable_tensor_generator(); + + // Serialize config and create loader + config.add_output("tensor_generator"); + std::string config_string = config.OutputAsString(); + DataLoader loader(config_string); + + return; + + std::atomic batch_count = 0; + auto start_time = absl::Now(); + + // Start logging thread + std::atomic should_stop{false}; + std::thread logging_thread([&]() { + while (!should_stop.load()) { + std::this_thread::sleep_for(std::chrono::seconds(1)); + + auto current_time = absl::Now(); + auto total_elapsed = current_time - start_time; + double rate = batch_count / absl::ToDoubleSeconds(total_elapsed); + + auto [stats_string, duration] = + loader.GetBucketMetrics(0, false); // k1Second = 0 + DataLoaderMetricsProto metrics; + metrics.ParseFromString(stats_string); + std::string metrics_json = metrics.OutputAsJson(); + + LOG(INFO) << absl::StrCat("Processed ", batch_count.load(), + " batches in ", + absl::ToDoubleSeconds(total_elapsed), + "s. Rate: ", absl::StrFormat("%.2f", rate), + " batches/sec. ", "Metrics: ", metrics_json); + } + }); + + while (true) { + TensorTuple batch = loader.GetNext("train"); + ++batch_count; + } +} + +} // namespace training +} // namespace lczero + +int main(int argc, char* argv[]) { + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo); + lczero::training::Run(); + return 0; +} diff --git a/csrc/loader/pybind_module.cc b/csrc/loader/pybind_module.cc new file mode 100644 index 00000000..378c2af2 --- /dev/null +++ b/csrc/loader/pybind_module.cc @@ -0,0 +1,194 @@ +// ABOUTME: PyBind11 binding module exposing C++ DataLoader to Python. +// ABOUTME: Handles configuration conversion and tensor memory management for +// numpy arrays. + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "loader/data_loader.h" +#include "loader/stages/chunk_source_loader.h" +#include "loader/stages/chunk_unpacker.h" +#include "loader/stages/file_path_provider.h" +#include "loader/stages/shuffling_chunk_pool.h" +#include "loader/stages/shuffling_frame_sampler.h" +#include "loader/stages/tensor_generator.h" +#include "utils/tensor.h" + +namespace py = pybind11; + +namespace lczero { +namespace training { + +namespace { + +std::string SerializePyProto(const py::handle& obj, const char* expected_type) { + if (py::isinstance(obj)) { + return obj.cast(); + } + if (py::hasattr(obj, "SerializeToString")) { + py::object bytes_obj = obj.attr("SerializeToString")(); + return bytes_obj.cast().cast(); + } + throw std::invalid_argument(std::string("Expected ") + expected_type + + " protobuf message or bytes."); +} + +template +ProtoT ParsePyProto(const py::handle& obj, const char* expected_type) { + ProtoT proto; + proto.ParseFromString(SerializePyProto(obj, expected_type)); + return proto; +} + +py::object MakePythonProto(const char* module_name, const char* message_name, + const std::string& serialized) { + py::object message_cls = py::module::import(module_name).attr(message_name); + py::object message_obj = message_cls(); + message_obj.attr("ParseFromString")(py::bytes(serialized)); + return message_obj; +} + +} // namespace + +// Helper function to convert TensorBase to numpy array using buffer protocol. +py::array tensor_to_numpy(std::unique_ptr tensor) { + // Extract raw pointer and release ownership from unique_ptr. + TensorBase* raw_tensor = tensor.release(); + + // Create numpy array with take_ownership policy. + // This transfers memory ownership to Python/numpy. + return py::array( + py::dtype(raw_tensor->py_format()), raw_tensor->shape(), + raw_tensor->strides(), raw_tensor->data(), + py::cast(raw_tensor, py::return_value_policy::take_ownership)); +} + +// Convert TensorTuple to tuple of numpy arrays. +py::tuple tensor_tuple_to_numpy_tuple(TensorTuple tensor_tuple) { + py::tuple result(tensor_tuple.size()); + for (size_t i = 0; i < tensor_tuple.size(); ++i) { + result[i] = tensor_to_numpy(std::move(tensor_tuple[i])); + } + return result; +} + +PYBIND11_MODULE(_lczero_training, m) { + absl::InitializeLog(); + absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo); + m.doc() = "Leela Chess Zero training data loader"; + + // Configuration is now handled via protobuf serialized strings + + // Expose the main DataLoader class. + py::class_(m, "DataLoader") + .def(py::init([](py::object config) { + std::string config_string = + SerializePyProto(config, "DataLoaderConfig"); + py::gil_scoped_release release; + return new DataLoader(config_string); + }), + py::arg("config"), + "Create DataLoader from DataLoaderConfig proto or bytes.") + .def( + "add_stages", + [](DataLoader& self, py::object config) { + std::string config_string = + SerializePyProto(config, "DataLoaderConfig"); + py::gil_scoped_release release; + self.AddStages(config_string); + }, + py::arg("config"), + "Append stages from DataLoaderConfig proto or bytes.") + .def( + "send_control_message", + [](DataLoader& self, py::object request) { + StageControlRequest control_request = + ParsePyProto(request, + "StageControlRequest"); + auto responses = [&]() { + py::gil_scoped_release release; + return self.SendControlMessage(control_request); + }(); + py::list result; + for (const auto& [stage_name, response] : responses) { + py::object response_obj = MakePythonProto( + "proto.stage_control_pb2", "StageControlResponse", + response.OutputAsString()); + result.append(py::make_tuple(stage_name, response_obj)); + } + return result; + }, + py::arg("request"), + "Send StageControlRequest to stages and return (stage_name, " + "StageControlResponse) tuples.") + .def( + "get_next", + [](DataLoader& self, const std::string& alias) { + return tensor_tuple_to_numpy_tuple([&] { + py::gil_scoped_release release; + return self.GetNext(alias); + }()); + }, + py::arg("alias") = "", + "Get next batch for the given output alias (default empty) as a " + "tuple of numpy arrays") + .def( + "maybe_get_next", + [](DataLoader& self, + const std::string& alias) -> std::optional { + auto result = [&] { + py::gil_scoped_release release; + return self.MaybeGetNext(alias); + }(); + if (result.has_value()) { + return tensor_tuple_to_numpy_tuple(std::move(*result)); + } + return std::nullopt; + }, + py::arg("alias") = "", + "Non-blocking get next batch for the given output alias (default " + "empty). Returns tuple of numpy arrays or None if no data available") + .def( + "get_bucket_metrics", + [](const DataLoader& self, int time_period, bool include_pending) { + auto [metrics, duration] = [&] { + py::gil_scoped_release release; + return self.GetBucketMetrics(time_period, include_pending); + }(); + return py::make_tuple(py::bytes(metrics), duration); + }, + "Get serialized metrics for bucket and duration as (bytes, float)") + .def( + "get_aggregate_ending_now", + [](const DataLoader& self, float duration_seconds, + bool include_pending) { + auto [metrics, duration] = [&] { + py::gil_scoped_release release; + return self.GetAggregateEndingNow(duration_seconds, + include_pending); + }(); + return py::make_tuple(py::bytes(metrics), duration); + }, + "Get serialized metrics for aggregate duration and actual duration " + "as (bytes, float)") + .def("start", &DataLoader::Start, "Start the data loader processing") + .def("stop", &DataLoader::Stop, "Stop the data loader"); + + // Expose TensorBase for potential advanced usage. + py::class_(m, "TensorBase") + .def("shape", &TensorBase::shape, py::return_value_policy::reference) + .def("strides", &TensorBase::strides, py::return_value_policy::reference) + .def("element_size", &TensorBase::element_size) + .def("py_format", &TensorBase::py_format); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_rescorer.cc b/csrc/loader/stages/chunk_rescorer.cc new file mode 100644 index 00000000..fcf31b6a --- /dev/null +++ b/csrc/loader/stages/chunk_rescorer.cc @@ -0,0 +1,177 @@ +#include "loader/stages/chunk_rescorer.h" + +#include +#include + +#include "absl/base/call_once.h" +#include "absl/log/log.h" +#include "chess/board.h" +#include "loader/data_loader_metrics.h" + +namespace lczero { +namespace training { +namespace { + +void V6ToV7(std::span data, float theta = 5.0f / 6.0f) { + if (data.empty()) return; + + const float beta = 1.0f - theta; + + float st_q = 0.0f; + float st_d = 0.0f; + + // Iterate backwards to calculate EMA and lookaheads in one go. + for (size_t i = data.size(); i-- > 0;) { + FrameType& item = data[i]; + + // For Q, we operate in an alternating sign domain. + const float sign = (i % 2 != 0) ? -1.0f : 1.0f; + const float cur_q = item.root_q * sign; + + if (i == data.size() - 1) { + st_q = cur_q; + st_d = item.root_d; + } else { + st_q = theta * st_q + beta * cur_q; + st_d = theta * st_d + beta * item.root_d; + } + + item.version = 7; + item.q_st = st_q * sign; + item.d_st = std::max(st_d, 0.0f); + + // Handle lookahead indices safely. + auto get_idx = [&](size_t offset) { + return (i + offset < data.size()) ? data[i + offset].played_idx : 65535; + }; + item.opp_played_idx = get_idx(1); + item.next_played_idx = get_idx(2); + } +} + +} // namespace + +ChunkRescorer::ChunkRescorer(const ChunkRescorerConfig& config) + : SingleInputStage(config), + SingleOutputStage(config.output()), + syzygy_paths_(config.syzygy_paths()), + dist_temp_(config.dist_temp()), + dist_offset_(config.dist_offset()), + dtz_boost_(config.dtz_boost()), + new_input_format_(config.new_input_format()), + thread_pool_(config.threads(), ThreadPoolOptions{}), + st_q_theta_(config.st_q_theta()) { + static absl::once_flag bitboards_initialized_flag; + absl::call_once(bitboards_initialized_flag, InitializeMagicBitboards); + + if (config.has_deblunder_threshold() && config.has_deblunder_width()) { + RescorerDeblunderSetup(config.deblunder_threshold(), + config.deblunder_width()); + } + + if (config.has_gaviota_paths()) { + RescorerGaviotaSetup(std::string(config.gaviota_paths())); + } + + LOG(INFO) << "Initializing ChunkRescorer with " << config.threads() + << " worker thread(s)"; + + thread_contexts_.reserve(config.threads()); + for (size_t i = 0; i < config.threads(); ++i) { + thread_contexts_.push_back(std::make_unique()); + } +} + +ChunkRescorer::~ChunkRescorer() { Stop(); } + +void ChunkRescorer::InitializeTablebase() { + if (tablebase_initialized_) { + return; + } + + LOG(INFO) << "ChunkRescorer initializing Syzygy tablebase with paths '" + << syzygy_paths_ << "'."; + tablebase_initialized_ = tablebase_.init(syzygy_paths_); + if (tablebase_initialized_) { + LOG(INFO) << "ChunkRescorer Syzygy max cardinality: " + << tablebase_.max_cardinality(); + } else { + LOG(WARNING) << "ChunkRescorer failed to initialize Syzygy tablebase; " + "rescoring will continue without tablebase lookups."; + } +} + +void ChunkRescorer::Start() { + LOG(INFO) << "Starting ChunkRescorer worker threads."; + InitializeTablebase(); + for (size_t i = 0; i < thread_contexts_.size(); ++i) { + thread_pool_.Enqueue([this, i](std::stop_token stop_token) { + Worker(stop_token, thread_contexts_[i].get()); + }); + } +} + +void ChunkRescorer::Stop() { + if (thread_pool_.stop_token().stop_requested()) return; + + LOG(INFO) << "Stopping ChunkRescorer."; + thread_pool_.Shutdown(); + output_queue()->Close(); + LOG(INFO) << "ChunkRescorer stopped."; +} + +void ChunkRescorer::Worker(std::stop_token stop_token, ThreadContext* context) { + auto producer = output_queue()->CreateProducer(); + + try { + while (true) { + TrainingChunk chunk = [&]() { + LoadMetricPauser pauser(context->load_metric_updater); + return input_queue()->Get(stop_token); + }(); + + try { + chunk.frames = RescoreTrainingData( + chunk.frames, &tablebase_, dist_temp_, dist_offset_, dtz_boost_, + new_input_format_); + if (chunk.frames.at(0).version == 6) V6ToV7(chunk.frames, st_q_theta_); + LoadMetricPauser pauser(context->load_metric_updater); + producer.Put(std::move(chunk), stop_token); + } catch (const std::exception& exception) { + failed_rescores_.fetch_add(1, std::memory_order_acq_rel); + LOG(ERROR) << "ChunkRescorer failed to rescore chunk: " + << exception.what() << "; sort_key=" << chunk.sort_key + << "; index_within_sort_key=" << chunk.index_within_sort_key + << "; global_index=" << chunk.global_index + << "; use_count=" << chunk.use_count + << "; frame_count=" << chunk.frames.size(); + continue; + } + } + } catch (const QueueClosedException&) { + LOG(INFO) << "ChunkRescorer worker stopping, queue closed."; + } catch (const QueueRequestCancelled&) { + LOG(INFO) << "ChunkRescorer worker stopping, request cancelled."; + } +} + +StageMetricProto ChunkRescorer::FlushMetrics() { + StageMetricProto stage_metric; + LoadMetricProto aggregated_load; + aggregated_load.set_name("load"); + for (const auto& context : thread_contexts_) { + UpdateFrom(aggregated_load, context->load_metric_updater.FlushMetrics()); + } + *stage_metric.add_load_metrics() = std::move(aggregated_load); + + auto* failed = stage_metric.add_count_metrics(); + failed->set_name("failed_rescores"); + failed->set_count(failed_rescores_.exchange(0, std::memory_order_acq_rel)); + + *stage_metric.add_queue_metrics() = + MetricsFromQueue("output", *output_queue()); + return stage_metric; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_rescorer.h b/csrc/loader/stages/chunk_rescorer.h new file mode 100644 index 00000000..be8255ad --- /dev/null +++ b/csrc/loader/stages/chunk_rescorer.h @@ -0,0 +1,65 @@ +// ABOUTME: Stage that rescales training chunks using Syzygy tablebases. +// ABOUTME: Adjusts frame metadata by invoking the classic LCZero rescorer. +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "libs/lc0/src/syzygy/syzygy.h" +#include "libs/lc0/src/trainingdata/rescorer.h" +#include "loader/frame_type.h" +#include "loader/stages/stage.h" +#include "loader/stages/training_chunk.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/metrics/load_metric.h" +#include "utils/queue.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +// Stage that takes TrainingChunk objects, applies tablebase-based rescoring and +// forwards the updated chunks downstream. +class ChunkRescorer + : public SingleInputStage, + public SingleOutputStage { + public: + using InputType = TrainingChunk; + using OutputType = TrainingChunk; + + explicit ChunkRescorer(const ChunkRescorerConfig& config); + ~ChunkRescorer() override; + + void Start() override; + void Stop() override; + StageMetricProto FlushMetrics() override; + + private: + struct ThreadContext { + LoadMetricUpdater load_metric_updater; + }; + + void Worker(std::stop_token stop_token, ThreadContext* context); + void InitializeTablebase(); + + SyzygyTablebase tablebase_; + bool tablebase_initialized_ = false; + std::string syzygy_paths_; + float dist_temp_; + float dist_offset_; + float dtz_boost_; + int new_input_format_; + + ThreadPool thread_pool_; + std::vector> thread_contexts_; + std::atomic failed_rescores_{0}; + float st_q_theta_; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_rescorer_test.cc b/csrc/loader/stages/chunk_rescorer_test.cc new file mode 100644 index 00000000..bbd60a70 --- /dev/null +++ b/csrc/loader/stages/chunk_rescorer_test.cc @@ -0,0 +1,82 @@ +#include "loader/stages/chunk_rescorer.h" + +#include +#include +#include + +#include "gtest/gtest.h" +#include "loader/stages/training_chunk.h" +#include "proto/data_loader_config.pb.h" +#include "utils/queue.h" + +namespace lczero { +namespace training { + +namespace { + +template +class PassthroughStage : public Stage { + public: + explicit PassthroughStage(Queue* queue) : queue_(queue) {} + + void Start() override {} + void Stop() override {} + StageMetricProto FlushMetrics() override { return StageMetricProto(); } + QueueBase* GetOutput(std::string_view name = "") override { + (void)name; + return queue_; + } + void SetInputs(absl::Span inputs) override { + if (!inputs.empty()) { + throw std::runtime_error("PassthroughStage expects no inputs"); + } + } + + private: + Queue* queue_; +}; + +} // namespace + +class ChunkRescorerTest : public ::testing::Test { + protected: + void SetUp() override { + input_queue_ = std::make_unique>(10); + config_.set_threads(1); + config_.mutable_output()->set_queue_capacity(10); + config_.set_syzygy_paths(""); + config_.set_dist_temp(0.75f); + config_.set_dist_offset(0.1f); + config_.set_dtz_boost(0.2f); + config_.set_new_input_format(-1); + } + + TrainingChunk MakeChunk(std::vector frames, + std::string sort_key = "alpha", size_t index = 3, + uint32_t use = 7) { + TrainingChunk chunk; + chunk.sort_key = std::move(sort_key); + chunk.index_within_sort_key = index; + chunk.use_count = use; + chunk.frames = std::move(frames); + return chunk; + } + + std::unique_ptr> input_queue_; + ChunkRescorerConfig config_; +}; + +TEST_F(ChunkRescorerTest, HandlesInputQueueClosure) { + ChunkRescorer rescorer(config_); + rescorer.SetInputs({input_queue_.get()}); + rescorer.Start(); + + input_queue_->Close(); + + EXPECT_THROW(rescorer.output_queue()->Get(), QueueClosedException); + + rescorer.Stop(); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_source_loader.cc b/csrc/loader/stages/chunk_source_loader.cc new file mode 100644 index 00000000..cb549dcf --- /dev/null +++ b/csrc/loader/stages/chunk_source_loader.cc @@ -0,0 +1,190 @@ +#include "loader/stages/chunk_source_loader.h" + +#include +#include + +#include "absl/log/log.h" +#include "loader/chunk_source/rawfile_chunk_source.h" +#include "loader/chunk_source/tar_chunk_source.h" +#include "loader/data_loader_metrics.h" +#include "proto/data_loader_config.pb.h" + +namespace lczero { +namespace training { + +std::unique_ptr CreateChunkSourceFromFile( + const std::filesystem::path& filepath, + ChunkSourceLoaderConfig::FrameFormat frame_format) { + auto extension = filepath.extension(); + try { + if (extension == ".gz") { + return std::make_unique(filepath, frame_format); + } + if (extension == ".tar") { + return std::make_unique(filepath, frame_format); + } + } catch (const std::exception& e) { + LOG(ERROR) << "Failed to create chunk source for " << filepath << ": " + << e.what(); + return nullptr; + } + return nullptr; +} + +ChunkSourceLoader::ChunkSourceLoader(const ChunkSourceLoaderConfig& config) + : SingleInputStage(config), + SingleOutputStage(config.output()), + thread_pool_(config.threads(), ThreadPoolOptions{}), + frame_format_(config.frame_format()) { + LOG(INFO) << "Initializing ChunkSourceLoader with " << config.threads() + << " worker threads"; + + // Initialize thread contexts but don't start worker threads yet. + thread_contexts_.reserve(config.threads()); + for (size_t i = 0; i < config.threads(); ++i) { + thread_contexts_.push_back(std::make_unique()); + } +} + +ChunkSourceLoader::~ChunkSourceLoader() { Stop(); } + +void ChunkSourceLoader::Start() { + LOG(INFO) << "Starting ChunkSourceLoader worker threads."; + for (size_t i = 0; i < thread_contexts_.size(); ++i) { + thread_pool_.Enqueue([this, i](std::stop_token stop_token) { + Worker(stop_token, thread_contexts_[i].get()); + }); + } +} + +void ChunkSourceLoader::Stop() { + if (thread_pool_.stop_token().stop_requested()) return; + + LOG(INFO) << "Stopping ChunkSourceLoader."; + thread_pool_.Shutdown(); + output_queue()->Close(); + LOG(INFO) << "ChunkSourceLoader stopped."; +} + +void ChunkSourceLoader::Worker(std::stop_token stop_token, + ThreadContext* context) { + auto producer = output_queue()->CreateProducer(); + LOG(INFO) << "ChunkSourceLoader worker@" << static_cast(context) + << " started."; + + try { + while (true) { + auto file = [&]() { + LoadMetricPauser pauser(context->load_metric_updater); + return input_queue()->Get(stop_token); + }(); + + if (file.message_type == + FilePathProvider::MessageType::kInitialScanComplete) { + LOG(INFO) + << "ChunkSourceLoader received initial scan completion marker."; + + bool should_forward; + { + absl::MutexLock lock(&phase_mutex_); + sentinel_received_ = true; + should_forward = (pre_sentinel_work_count_ == 0); + } + + if (should_forward) { + LOG(INFO) << "ChunkSourceLoader forwarding initial scan completion " + "marker."; + producer.Put({.source = nullptr, .message_type = file.message_type}, + stop_token); + } + continue; + } + + // Track pre-sentinel work. + bool is_pre_sentinel; + { + absl::MutexLock lock(&phase_mutex_); + is_pre_sentinel = !sentinel_received_; + if (is_pre_sentinel) pre_sentinel_work_count_++; + } + + // Create ChunkSource from the file. + LOG_EVERY_N(INFO, 1000) + << "ChunkSourceLoader preparing chunk source for " << file.filepath; + auto source = CreateChunkSourceFromFile(file.filepath, frame_format_); + if (source) { + { + absl::MutexLock lock(&last_chunk_key_mutex_); + last_chunk_key_ = source->GetChunkSortKey(); + } + ChunkSourceWithPhase output{.source = std::move(source), + .message_type = file.message_type}; + LoadMetricPauser pauser(context->load_metric_updater); + producer.Put(std::move(output), stop_token); + } else { + LOG_EVERY_N(INFO, 100) + << "ChunkSourceLoader skipping unsupported file: " << file.filepath; + skipped_files_count_++; + } + + // Complete pre-sentinel work tracking. + if (is_pre_sentinel) { + absl::MutexLock lock(&phase_mutex_); + if (--pre_sentinel_work_count_ == 0 && sentinel_received_) { + LOG(INFO) << "ChunkSourceLoader forwarding initial scan completion " + "marker after all pre-sentinel work completed."; + producer.Put( + {.source = nullptr, + .message_type = + FilePathProvider::MessageType::kInitialScanComplete}, + stop_token); + } + } + } + } catch (const QueueClosedException&) { + LOG(INFO) << "ChunkSourceLoader worker@" + << static_cast(context) + << " stopping, queue closed."; + } catch (const QueueRequestCancelled&) { + LOG(INFO) << "ChunkSourceLoader worker@" + << static_cast(context) + << " stopping, request cancelled."; + } catch (const std::exception& e) { + LOG(ERROR) << "ChunkSourceLoader worker@" + << static_cast(context) + << " exiting due to exception: " << e.what(); + throw; + } + + LOG(INFO) << "ChunkSourceLoader worker@" << static_cast(context) + << " exiting loop."; +} + +StageMetricProto ChunkSourceLoader::FlushMetrics() { + StageMetricProto stage_metric; + LoadMetricProto aggregated_load; + aggregated_load.set_name("load"); + for (const auto& context : thread_contexts_) { + UpdateFrom(aggregated_load, context->load_metric_updater.FlushMetrics()); + } + *stage_metric.add_load_metrics() = std::move(aggregated_load); + + auto* skipped_metric = stage_metric.add_count_metrics(); + skipped_metric->set_name("skipped_files"); + skipped_metric->set_count(skipped_files_count_.exchange(0)); + + // Get the last chunk key. + { + absl::MutexLock lock(&last_chunk_key_mutex_); + if (!last_chunk_key_.empty()) { + stage_metric.set_last_chunk_key(last_chunk_key_); + } + } + + *stage_metric.add_queue_metrics() = + MetricsFromQueue("output", *output_queue()); + return stage_metric; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_source_loader.h b/csrc/loader/stages/chunk_source_loader.h new file mode 100644 index 00000000..c7b66e29 --- /dev/null +++ b/csrc/loader/stages/chunk_source_loader.h @@ -0,0 +1,70 @@ +#pragma once + +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/synchronization/mutex.h" +#include "loader/chunk_source/chunk_source.h" +#include "loader/stages/file_path_provider.h" +#include "loader/stages/stage.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/metrics/load_metric.h" +#include "utils/queue.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +// Creates a ChunkSource based on file extension. Returns RawFileChunkSource for +// .gz files, TarChunkSource for .tar files, or nullptr for unsupported types. +std::unique_ptr CreateChunkSourceFromFile( + const std::filesystem::path& filepath, + ChunkSourceLoaderConfig::FrameFormat frame_format); + +struct ChunkSourceWithPhase { + std::unique_ptr source; + FilePathProvider::MessageType message_type; +}; + +// Worker pool that converts FilePathProvider output to ChunkSource objects. +// Takes FilePathProvider::File as input and outputs ChunkSourceWithPhase. +class ChunkSourceLoader + : public SingleInputStage, + public SingleOutputStage { + public: + using InputType = FilePathProvider::File; + using OutputType = ChunkSourceWithPhase; + + explicit ChunkSourceLoader(const ChunkSourceLoaderConfig& config); + ~ChunkSourceLoader(); + + void Start() override; + void Stop() override; + + StageMetricProto FlushMetrics() override; + + private: + struct ThreadContext { + LoadMetricUpdater load_metric_updater; + }; + + void Worker(std::stop_token stop_token, ThreadContext* context); + ThreadPool thread_pool_; + std::vector> thread_contexts_; + std::atomic skipped_files_count_{0}; + absl::Mutex last_chunk_key_mutex_; + std::string last_chunk_key_; + ChunkSourceLoaderConfig::FrameFormat frame_format_; + + // Synchronization for sentinel barrier. + absl::Mutex phase_mutex_; + int pre_sentinel_work_count_ ABSL_GUARDED_BY(phase_mutex_) = 0; + bool sentinel_received_ ABSL_GUARDED_BY(phase_mutex_) = false; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_source_loader_test.cc b/csrc/loader/stages/chunk_source_loader_test.cc new file mode 100644 index 00000000..08d3a153 --- /dev/null +++ b/csrc/loader/stages/chunk_source_loader_test.cc @@ -0,0 +1,199 @@ +#include "loader/stages/chunk_source_loader.h" + +#include + +#include + +#include "loader/stages/file_path_provider.h" +#include "utils/queue.h" + +namespace lczero { +namespace training { + +namespace { + +template +class PassthroughStage : public Stage { + public: + explicit PassthroughStage(Queue* queue) : queue_(queue) {} + + void Start() override {} + void Stop() override {} + StageMetricProto FlushMetrics() override { return StageMetricProto(); } + QueueBase* GetOutput(std::string_view name = "") override { + (void)name; + return queue_; + } + void SetInputs(absl::Span inputs) override { + if (!inputs.empty()) { + throw std::runtime_error("PassthroughStage expects no inputs"); + } + } + + private: + Queue* queue_; +}; + +} // namespace + +TEST(ChunkSourceLoaderTest, ProcessesFiles) { + Queue input_queue(10); + ChunkSourceLoaderConfig config; + config.set_threads(1); + config.mutable_output()->set_queue_capacity(10); + ChunkSourceLoader feed(config); + feed.SetInputs({&input_queue}); + feed.Start(); + + { + auto producer = input_queue.CreateProducer(); + // Add a file with unsupported extension (should not create ChunkSource) + producer.Put(FilePathProvider::File{ + .filepath = + std::filesystem::path("/test.txt"), // unsupported extension + .message_type = FilePathProvider::MessageType::kFile}); + } // Producer destroyed here, closing input queue + + // Try to get output - there should be no valid ChunkSources for unsupported + // files + try { + while (true) { + auto output = feed.output_queue()->Get(); + // If we get output, it means a ChunkSource was created, which shouldn't + // happen for unsupported files + FAIL() << "Expected no output for unsupported file extension"; + } + } catch (const QueueClosedException&) { + // Expected: queue should be closed when input is done and no output + // produced + SUCCEED(); + } +} + +TEST(ChunkSourceLoaderTest, HandlesPhases) { + Queue input_queue(10); + ChunkSourceLoaderConfig config; + config.set_threads(1); + config.mutable_output()->set_queue_capacity(10); + ChunkSourceLoader feed(config); + feed.SetInputs({&input_queue}); + feed.Start(); + + { + auto producer = input_queue.CreateProducer(); + // Test different phases - all should be passed through even if no + // ChunkSource is created + producer.Put(FilePathProvider::File{ + .filepath = std::filesystem::path("/test1.gz"), + .message_type = FilePathProvider::MessageType::kFile}); + + producer.Put(FilePathProvider::File{ + .filepath = std::filesystem::path("/test2.gz"), + .message_type = FilePathProvider::MessageType::kFile}); + } // Producer destroyed here, closing input queue + + // Queue should eventually close when input is done + try { + while (true) { + feed.output_queue()->Get(); + } + } catch (const QueueClosedException&) { + SUCCEED(); + } +} + +TEST(ChunkSourceLoaderTest, PassesThroughInitialScanComplete) { + Queue input_queue(10); + ChunkSourceLoaderConfig config; + config.set_threads(1); + config.mutable_output()->set_queue_capacity(10); + ChunkSourceLoader feed(config); + feed.SetInputs({&input_queue}); + feed.Start(); + + { + auto producer = input_queue.CreateProducer(); + producer.Put(FilePathProvider::File{ + .filepath = std::filesystem::path(""), + .message_type = FilePathProvider::MessageType::kInitialScanComplete}); + } // Producer destroyed here, closing input queue + + // Should get kInitialScanComplete in output with null ChunkSource + auto output = feed.output_queue()->Get(); + EXPECT_EQ(output.message_type, + FilePathProvider::MessageType::kInitialScanComplete); + EXPECT_EQ(output.source, nullptr); + + // Queue should be closed after the single message + try { + feed.output_queue()->Get(); + FAIL() << "Expected queue to be closed"; + } catch (const QueueClosedException&) { + SUCCEED(); + } +} + +TEST(ChunkSourceLoaderTest, SentinelBarrierWithMultipleThreads) { + Queue input_queue(100); + ChunkSourceLoaderConfig config; + config.set_threads(4); + config.mutable_output()->set_queue_capacity(100); + ChunkSourceLoader feed(config); + feed.SetInputs({&input_queue}); + feed.Start(); + + { + auto producer = input_queue.CreateProducer(); + // Add files that will be processed before sentinel. + for (int i = 0; i < 20; ++i) { + producer.Put(FilePathProvider::File{ + .filepath = + std::filesystem::path("/test" + std::to_string(i) + ".txt"), + .message_type = FilePathProvider::MessageType::kFile}); + } + // Add sentinel. + producer.Put(FilePathProvider::File{ + .filepath = std::filesystem::path(""), + .message_type = FilePathProvider::MessageType::kInitialScanComplete}); + // Add files that arrive after sentinel. + for (int i = 20; i < 30; ++i) { + producer.Put(FilePathProvider::File{ + .filepath = + std::filesystem::path("/test" + std::to_string(i) + ".txt"), + .message_type = FilePathProvider::MessageType::kFile}); + } + } // Producer destroyed here, closing input queue + + // Read all outputs and verify sentinel comes after all pre-sentinel files. + int files_before_sentinel = 0; + int files_after_sentinel = 0; + bool sentinel_seen = false; + + try { + while (true) { + auto output = feed.output_queue()->Get(); + if (output.message_type == + FilePathProvider::MessageType::kInitialScanComplete) { + EXPECT_FALSE(sentinel_seen) << "Sentinel should appear exactly once"; + sentinel_seen = true; + } else { + if (sentinel_seen) { + files_after_sentinel++; + } else { + files_before_sentinel++; + } + } + } + } catch (const QueueClosedException&) { + } + + // Verify sentinel was seen. + EXPECT_TRUE(sentinel_seen); + // All 20 pre-sentinel files should be before sentinel (unsupported, so 0). + EXPECT_EQ(files_before_sentinel, 0); + // All 10 post-sentinel files should be after sentinel (unsupported, so 0). + EXPECT_EQ(files_after_sentinel, 0); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_source_splitter.cc b/csrc/loader/stages/chunk_source_splitter.cc new file mode 100644 index 00000000..cbe3e706 --- /dev/null +++ b/csrc/loader/stages/chunk_source_splitter.cc @@ -0,0 +1,168 @@ +#include "loader/stages/chunk_source_splitter.h" + +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/strings/str_cat.h" +#include "loader/data_loader_metrics.h" + +namespace lczero { +namespace training { + +ChunkSourceSplitter::ChunkSourceSplitter( + const ChunkSourceSplitterConfig& config) + : SingleInputStage(config) { + if (config.output().empty()) { + throw std::runtime_error("ChunkSourceSplitter requires at least 1 output."); + } + + // Validate parallel arrays have same size. + if (config.output_size() != config.weight_size()) { + throw std::runtime_error(absl::StrCat( + "ChunkSourceSplitter output and weight arrays must have same size: ", + config.output_size(), " vs ", config.weight_size())); + } + + // Create output queues from parallel arrays. + outputs_.reserve(config.output_size()); + for (size_t i = 0; i < static_cast(config.output_size()); ++i) { + const auto& queue_cfg = config.output(static_cast(i)); + const uint64_t weight = i < static_cast(config.weight_size()) + ? config.weight(static_cast(i)) + : 1; + + if (absl::c_any_of(outputs_, [&](const auto& existing_out) { + return existing_out->name == queue_cfg.name(); + })) { + throw std::runtime_error(std::string(absl::StrCat( + "Duplicate output name in ChunkSourceSplitter: ", queue_cfg.name()))); + } + + auto* out = outputs_ + .emplace_back(std::make_unique( + queue_cfg.name(), weight, queue_cfg.queue_capacity(), + ToOverflowBehavior(queue_cfg.overflow_behavior()))) + .get(); + LOG(INFO) << "ChunkSourceSplitter configured output '" << out->name + << "' weight=" << out->weight + << " capacity=" << queue_cfg.queue_capacity(); + } + + // Precompute cumulative weights for fast assignment. + cumulative_.resize(outputs_.size()); + std::transform_inclusive_scan( + outputs_.begin(), outputs_.end(), cumulative_.begin(), + std::plus{}, + [](const std::unique_ptr& out) { return out->weight; }); + + // Validate total weight is positive. + if (cumulative_.back() == 0) { + throw std::runtime_error( + "ChunkSourceSplitter requires at least one output with positive " + "weight."); + } +} + +ChunkSourceSplitter::~ChunkSourceSplitter() { Stop(); } + +void ChunkSourceSplitter::Start() { + LOG(INFO) << "Starting ChunkSourceSplitter worker."; + thread_pool_.Enqueue( + [this](std::stop_token stop_token) { Worker(stop_token); }); +} + +void ChunkSourceSplitter::Stop() { + if (thread_pool_.stop_token().stop_requested()) return; + LOG(INFO) << "Stopping ChunkSourceSplitter."; + thread_pool_.Shutdown(); + for (auto& out : outputs_) out->queue.Close(); +} + +QueueBase* ChunkSourceSplitter::GetOutput(std::string_view name) { + auto iter = absl::c_find_if( + outputs_, [&](const auto& out) { return out->name == name; }); + if (iter == outputs_.end()) { + throw std::runtime_error( + absl::StrCat("Unknown output '", name, "' for ChunkSourceSplitter.")); + } + return &(*iter)->queue; +} + +StageMetricProto ChunkSourceSplitter::FlushMetrics() { + StageMetricProto metric; + for (auto& out : outputs_) { + *metric.add_queue_metrics() = MetricsFromQueue(out->name, out->queue); + } + return metric; +} + +void ChunkSourceSplitter::Worker(std::stop_token stop_token) { + // Create producers for each output in this thread. + std::vector::Producer> producers; + producers.reserve(outputs_.size()); + for (auto& out : outputs_) { + producers.emplace_back(out->queue.CreateProducer()); + } + + try { + while (true) { + InputType item = input_queue()->Get(stop_token); + if (item.message_type == + FilePathProvider::MessageType::kInitialScanComplete) { + // Broadcast to all outputs. + for (auto& prod : producers) { + prod.Put( + OutputType{.source = nullptr, .message_type = item.message_type}, + stop_token); + } + continue; + } + + // Share ownership of the ChunkSource with any produced views. + std::shared_ptr shared_source(std::move(item.source)); + + auto per_output_indices = BuildAssignments(shared_source); + // Emit only non-empty views, preserving original message type (kFile). + for (size_t i = 0; i < outputs_.size(); ++i) { + if (per_output_indices[i].empty()) continue; + auto view = std::make_unique( + shared_source, std::move(per_output_indices[i])); + producers[i].Put( + OutputType{.source = std::move(view), + .message_type = FilePathProvider::MessageType::kFile}, + stop_token); + } + } + } catch (const QueueClosedException&) { + // Input queue closed — producers will close queues automatically when + // destroyed if this thread holds the last producer. + LOG(INFO) << "ChunkSourceSplitter worker exiting: input closed."; + } +} + +std::vector> ChunkSourceSplitter::BuildAssignments( + const std::shared_ptr& source) { + const std::string sort_key = source->GetChunkSortKey(); + const size_t n = source->GetChunkCount(); + + // Prepare result containers with a rough reservation. + std::vector> indices(outputs_.size()); + + for (size_t i = 0; i < n; ++i) { + const uint64_t h = + static_cast(absl::Hash>{}( + std::make_pair(sort_key, i))); + const uint64_t r = h % cumulative_.back(); + // Find the output where cumulative[j-1] <= r < cumulative[j]. + const auto it = std::upper_bound(cumulative_.begin(), cumulative_.end(), r); + const size_t idx = it - cumulative_.begin(); + indices[idx].push_back(static_cast(i)); + } + + return indices; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_source_splitter.h b/csrc/loader/stages/chunk_source_splitter.h new file mode 100644 index 00000000..156ef7a2 --- /dev/null +++ b/csrc/loader/stages/chunk_source_splitter.h @@ -0,0 +1,65 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/hash/hash.h" +#include "absl/log/log.h" +#include "loader/chunk_source/chunk_source.h" +#include "loader/chunk_source/chunk_source_view.h" +#include "loader/stages/chunk_source_loader.h" +#include "loader/stages/stage.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/queue.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +// Splits an incoming ChunkSource into several ChunkSourceViews based on a +// deterministic hash of (sort_key, index). Emits to multiple named outputs. +class ChunkSourceSplitter + : public SingleInputStage { + public: + using InputType = ChunkSourceWithPhase; + using OutputType = ChunkSourceWithPhase; + + explicit ChunkSourceSplitter(const ChunkSourceSplitterConfig& config); + ~ChunkSourceSplitter(); + + void Start() override; + void Stop() override; + + StageMetricProto FlushMetrics() override; + QueueBase* GetOutput(std::string_view name = "") override; + + private: + struct Output { + std::string name; + uint64_t weight; + Queue queue; + Output(std::string_view name, uint64_t weight, size_t capacity, + OverflowBehavior overflow) + : name(name), weight(weight), queue(capacity, overflow) {} + }; + + void Worker(std::stop_token stop_token); + + // Builds per-output indices given a source; uses absl::Hash on + // (sort_key, index) and weights to assign indices. + std::vector> BuildAssignments( + const std::shared_ptr& source); + + std::vector> outputs_; + std::vector cumulative_; + ThreadPool thread_pool_{1}; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_source_splitter_test.cc b/csrc/loader/stages/chunk_source_splitter_test.cc new file mode 100644 index 00000000..a3a4002a --- /dev/null +++ b/csrc/loader/stages/chunk_source_splitter_test.cc @@ -0,0 +1,162 @@ +#include "loader/stages/chunk_source_splitter.h" + +#include +#include +#include +#include + +#include "absl/hash/hash.h" +#include "gtest/gtest.h" +#include "loader/chunk_source/chunk_source.h" +#include "loader/stages/chunk_source_loader.h" +#include "loader/stages/stage.h" +#include "proto/data_loader_config.pb.h" +#include "utils/queue.h" + +namespace lczero { +namespace training { +namespace { + +// Simple fixed-count chunk source for testing. +class FixedCountChunkSource : public ChunkSource { + public: + FixedCountChunkSource(std::string sort_key, size_t count) + : key_(std::move(sort_key)), count_(count) {} + + private: + std::string GetChunkSortKey() const override { return key_; } + size_t GetChunkCount() const override { return count_; } + std::optional> GetChunkData(size_t) override { + return std::vector{FrameType{}}; + } + + std::string key_; + size_t count_; +}; + +template +class PassthroughStage : public Stage { + public: + explicit PassthroughStage(Queue* queue) : queue_(queue) {} + + void Start() override {} + void Stop() override {} + StageMetricProto FlushMetrics() override { return StageMetricProto(); } + QueueBase* GetOutput(std::string_view) override { return queue_; } + void SetInputs(absl::Span inputs) override { + if (!inputs.empty()) { + throw std::runtime_error("PassthroughStage expects no inputs"); + } + } + + private: + Queue* queue_; +}; + +} // namespace + +TEST(ChunkSourceSplitterTest, SplitsByHashAndWeight) { + // Upstream queue. + auto input_queue = std::make_unique>(8); + + // Configure splitter with two outputs A:1, B:2. + ChunkSourceSplitterConfig cfg; + auto* outA = cfg.add_output(); + outA->set_name("A"); + outA->set_queue_capacity(8); + cfg.add_weight(1); + auto* outB = cfg.add_output(); + outB->set_name("B"); + outB->set_queue_capacity(8); + cfg.add_weight(2); + + ChunkSourceSplitter splitter(cfg); + splitter.SetInputs({input_queue.get()}); + + splitter.Start(); + + // Send a source with known key and count. + const std::string key = "skey"; + const size_t count = 100; + ChunkSourceWithPhase item; + item.source = std::make_unique(key, count); + item.message_type = FilePathProvider::MessageType::kFile; + auto producer = input_queue->CreateProducer(); + producer.Put(std::move(item)); + producer.Close(); + + // Compute expected assignment counts using the same hash/weights. + const uint64_t total_weight = 3; + uint64_t cumA = 1; // [0] + size_t expectedA = 0; + size_t expectedB = 0; + for (size_t i = 0; i < count; ++i) { + const uint64_t h = static_cast( + absl::Hash>{}(std::make_pair(key, i))); + const uint64_t r = h % total_weight; + if (r < cumA) + ++expectedA; + else + ++expectedB; + } + + // Read outputs and verify view sizes. + auto* qa = + dynamic_cast*>(splitter.GetOutput("A")); + auto* qb = + dynamic_cast*>(splitter.GetOutput("B")); + + ASSERT_NE(qa, nullptr); + ASSERT_NE(qb, nullptr); + + auto msgA = qa->Get(); + auto msgB = qb->Get(); + ASSERT_NE(msgA.source, nullptr); + ASSERT_NE(msgB.source, nullptr); + EXPECT_EQ(msgA.source->GetChunkCount(), expectedA); + EXPECT_EQ(msgB.source->GetChunkCount(), expectedB); + + splitter.Stop(); +} + +TEST(ChunkSourceSplitterTest, BroadcastsInitialScanComplete) { + auto input_queue = std::make_unique>(4); + + ChunkSourceSplitterConfig cfg; + auto* outA = cfg.add_output(); + outA->set_name("A"); + cfg.add_weight(1); + auto* outB = cfg.add_output(); + outB->set_name("B"); + cfg.add_weight(1); + + ChunkSourceSplitter splitter(cfg); + splitter.SetInputs({input_queue.get()}); + splitter.Start(); + + ChunkSourceWithPhase marker; + marker.source = nullptr; + marker.message_type = FilePathProvider::MessageType::kInitialScanComplete; + auto producer = input_queue->CreateProducer(); + producer.Put(std::move(marker)); + producer.Close(); + + auto* qa = + dynamic_cast*>(splitter.GetOutput("A")); + auto* qb = + dynamic_cast*>(splitter.GetOutput("B")); + + auto m1 = qa->Get(); + auto m2 = qb->Get(); + EXPECT_EQ(m1.message_type, + FilePathProvider::MessageType::kInitialScanComplete); + EXPECT_EQ(m2.message_type, + FilePathProvider::MessageType::kInitialScanComplete); + EXPECT_EQ(m1.source, nullptr); + EXPECT_EQ(m2.source, nullptr); + + splitter.Stop(); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_unpacker.cc b/csrc/loader/stages/chunk_unpacker.cc new file mode 100644 index 00000000..9f754072 --- /dev/null +++ b/csrc/loader/stages/chunk_unpacker.cc @@ -0,0 +1,289 @@ +#include "loader/stages/chunk_unpacker.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "absl/numeric/int128.h" +#include "absl/random/bit_gen_ref.h" +#include "absl/random/random.h" +#include "loader/data_loader_metrics.h" +#include "loader/stages/position_sampling.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" + +namespace lczero { +namespace training { + +// Deterministically partitions `n` positions into disjoint subsets of size +// `~p*n`, returning the subset for a given `iteration`. While each selection +// individually behaves like a Bernoulli sample with probability `p`, the +// samples are correlated to ensure all positions are selected exactly once over +// `1/p` iterations. To sample disjoint subsets from the same set of positions, +// `gen` must be seeded identically for each call with a different `iteration`. +std::vector PickSampledPositions(int32_t n, double p, + int32_t iteration, + absl::BitGen& gen) { + assert(p > 0.0 && p <= 1.0); + double carried_prob = p; + + std::vector result; + absl::flat_hash_set skip_next_round; + + while (true) { + int32_t num_this_round = (1.0 - carried_prob) / p + 1; + double last_partial_prob = 1 - (carried_prob + (num_this_round - 1) * p); + const bool return_this_round = iteration < num_this_round; + absl::flat_hash_set skip_this_round(skip_next_round); + skip_next_round.clear(); + for (int32_t i = 0; i < n; ++i) { + if (skip_this_round.contains(i)) continue; + const double toss = absl::Uniform(gen, 0.0, 1.0); + const int32_t value = (toss - carried_prob) / p + 1; + if (value == iteration) result.push_back(static_cast(i)); + if (value >= num_this_round) { + skip_next_round.insert(static_cast(i)); + } + } + if (return_this_round) return result; + iteration -= num_this_round; + carried_prob = p - last_partial_prob; + } +} + +// Samples `k` indices from an infinite, deterministically-generated sequence, +// after an initial `skip`. The sequence is formed by repeatedly shuffling +// `[0..n-1]` and probabilistically selecting each index. Determinism is +// achieved by seeding a new PRNG for each shuffled block from a stable +// `root_seed` and the block's index. +std::vector SampleProbabilisticSequence( + uint64_t k, uint64_t skip, std::span probabilities, + absl::BitGen& gen) { + const size_t n = probabilities.size(); + if (n == 0 || k == 0) return {}; + + uint64_t skipped_so_far = 0; + std::vector v(n); + std::iota(v.begin(), v.end(), 0u); + std::vector result; + result.reserve(k); + + while (true) { + std::shuffle(v.begin(), v.end(), gen); + for (const uint32_t candidate : v) { + if (!absl::Bernoulli(gen, probabilities[candidate])) continue; + if (skipped_so_far < skip) { + skipped_so_far++; + } else { + result.push_back(candidate); + if (result.size() == k) return result; + } + } + } +} + +namespace { + +uint32_t GenerateRunSeed() { + absl::BitGen gen(absl::MakeSeedSeq()); + return absl::Uniform(gen); +} + +} // namespace + +ChunkUnpacker::ChunkUnpacker(const ChunkUnpackerConfig& config) + : SingleInputStage(config), + config_(config), + run_seed_(GenerateRunSeed()), + primary_output_queue_( + config.output().queue_capacity(), + ToOverflowBehavior(config.output().overflow_behavior())), + thread_pool_(config.threads(), ThreadPoolOptions{}) { + const bool has_rate = config.has_position_sampling_rate(); + const bool has_count = config.has_position_count(); + const bool has_prefetch_count = config.has_prefetch_count(); + const bool has_prefetch_output = config.has_prefetch_output(); + + CHECK(has_prefetch_count == has_prefetch_output) + << "prefetch_count and prefetch_output must both be set or both unset."; + + if (has_prefetch_count) { + CHECK(has_count) << "position_count must be set when using prefetch mode."; + CHECK(!has_rate) + << "position_sampling_rate cannot be used in prefetch mode."; + CHECK(config.position_count() == 1) + << "position_count must equal 1 in prefetch mode, got " + << config.position_count(); + prefetch_output_queue_.emplace( + config.prefetch_output().queue_capacity(), + ToOverflowBehavior(config.prefetch_output().overflow_behavior())); + if (config.output().name() == config.prefetch_output().name()) { + throw std::runtime_error( + absl::StrCat("ChunkUnpacker output names must be different, got: '", + config.output().name(), "'")); + } + } else { + CHECK(has_rate != has_count) + << "Exactly one of position_sampling_rate or position_count must be " + "set."; + } + + LOG(INFO) << "Initializing ChunkUnpacker with " << config.threads() + << " worker threads"; + + // Initialize thread contexts but don't start worker threads yet. + thread_contexts_.reserve(config.threads()); + for (size_t i = 0; i < config.threads(); ++i) { + thread_contexts_.push_back(std::make_unique()); + } +} + +ChunkUnpacker::~ChunkUnpacker() { Stop(); } + +void ChunkUnpacker::Start() { + LOG(INFO) << "Starting ChunkUnpacker worker threads."; + for (size_t i = 0; i < thread_contexts_.size(); ++i) { + thread_pool_.Enqueue([this, i](std::stop_token stop_token) { + Worker(stop_token, thread_contexts_[i].get()); + }); + } +} + +void ChunkUnpacker::Stop() { + if (thread_pool_.stop_token().stop_requested()) return; + + LOG(INFO) << "Stopping ChunkUnpacker."; + thread_pool_.Shutdown(); + primary_output_queue_.Close(); + if (prefetch_output_queue_) prefetch_output_queue_->Close(); + LOG(INFO) << "ChunkUnpacker stopped."; +} + +QueueBase* ChunkUnpacker::GetOutput(std::string_view name) { + if (name == config_.output().name()) return &primary_output_queue_; + if (config_.has_prefetch_output() && + name == config_.prefetch_output().name()) { + return &*prefetch_output_queue_; + } + std::string available = absl::StrCat("'", config_.output().name(), "'"); + if (config_.has_prefetch_output()) { + absl::StrAppend(&available, ", '", config_.prefetch_output().name(), "'"); + } + throw std::runtime_error(absl::StrCat("ChunkUnpacker unknown output '", name, + "'. Available outputs: ", available)); +} + +namespace { +std::vector FramesToProbabilities(std::span frames, + const PositionSamplingConfig& config) { + std::vector probabilities; + probabilities.reserve(frames.size()); + absl::c_transform(frames, std::back_inserter(probabilities), + [&](const FrameType& frame) { + return ComputePositionSamplingWeight(frame, config); + }); + const float max_prob = *absl::c_max_element(probabilities); + if (max_prob > 0.0f) { + absl::c_transform(probabilities, probabilities.begin(), + [max_prob](float p) { return p / max_prob; }); + } else { + absl::c_fill(probabilities, 1.0f); + } + return probabilities; +} +} // namespace + +void ChunkUnpacker::Worker(std::stop_token stop_token, ThreadContext* context) { + // Create a local producer for this worker thread. + auto primary_producer = primary_output_queue_.CreateProducer(); + std::optionalCreateProducer())> + prefetch_producer; + if (prefetch_output_queue_.has_value()) { + prefetch_producer.emplace(prefetch_output_queue_->CreateProducer()); + } + + try { + while (true) { + auto chunk = [&]() { + LoadMetricPauser pauser(context->load_metric_updater); + return input_queue()->Get(stop_token); + }(); + + absl::BitGen gen( + std::seed_seq{run_seed_, static_cast(chunk.global_index)}); + std::vector positions; + if (config_.has_position_sampling_rate()) { + positions = PickSampledPositions( + static_cast(chunk.frames.size()), + config_.position_sampling_rate(), chunk.use_count, gen); + } else { + auto probabilities = + FramesToProbabilities(chunk.frames, config_.position_sampling()); + positions = SampleProbabilisticSequence( + config_.position_count() + config_.prefetch_count(), + config_.position_count() * chunk.use_count, probabilities, gen); + } + + if (config_.has_prefetch_count()) { + // Prefetch mode: output first position to primary, rest to prefetch. + if (!positions.empty()) { + LoadMetricPauser pauser(context->load_metric_updater); + primary_producer.Put(std::move(chunk.frames[positions[0]]), + stop_token); + } + if (positions.size() > 1) { + CacheRequest cache_request; + cache_request.global_index = chunk.global_index; + cache_request.next_use = chunk.use_count + 1; + cache_request.items.reserve(positions.size() - 1); + for (size_t i = 1; i < positions.size(); ++i) { + cache_request.items.push_back(chunk.frames[positions[i]]); + } + LoadMetricPauser pauser(context->load_metric_updater); + prefetch_producer->Put(std::move(cache_request), stop_token); + } + } else { + // Normal mode: output all positions to primary. + for (uint32_t pos : positions) { + LoadMetricPauser pauser(context->load_metric_updater); + primary_producer.Put(std::move(chunk.frames[pos]), stop_token); + } + } + } + } catch (const QueueClosedException&) { + LOG(INFO) << "ChunkUnpacker worker stopping, queue closed."; + } catch (const QueueRequestCancelled&) { + LOG(INFO) << "ChunkUnpacker worker stopping, request cancelled."; + } +} + +StageMetricProto ChunkUnpacker::FlushMetrics() { + StageMetricProto stage_metric; + LoadMetricProto aggregated_load; + aggregated_load.set_name("load"); + for (const auto& context : thread_contexts_) { + UpdateFrom(aggregated_load, context->load_metric_updater.FlushMetrics()); + } + *stage_metric.add_load_metrics() = std::move(aggregated_load); + *stage_metric.add_queue_metrics() = + MetricsFromQueue(config_.output().name(), primary_output_queue_); + if (prefetch_output_queue_.has_value()) { + *stage_metric.add_queue_metrics() = MetricsFromQueue( + config_.prefetch_output().name(), *prefetch_output_queue_); + } + return stage_metric; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_unpacker.h b/csrc/loader/stages/chunk_unpacker.h new file mode 100644 index 00000000..250c12b3 --- /dev/null +++ b/csrc/loader/stages/chunk_unpacker.h @@ -0,0 +1,65 @@ +// ABOUTME: Stage that unpacks chunks into FrameType frames. +// ABOUTME: Converts stream of std::string chunks to FrameType stream. +#pragma once + +#include +#include +#include +#include +#include + +#include "absl/random/random.h" +#include "absl/types/optional.h" +#include "loader/data_loader_metrics.h" +#include "loader/frame_type.h" +#include "loader/stages/stage.h" +#include "loader/stages/training_chunk.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/queue.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +// Worker pool that unpacks chunks into frames. +// Takes parsed TrainingChunk objects as input and outputs individual +// FrameType frames. +class ChunkUnpacker + : public SingleInputStage { + public: + using InputType = TrainingChunk; + + explicit ChunkUnpacker(const ChunkUnpackerConfig& config); + ~ChunkUnpacker(); + + void Start() override; + void Stop() override; + QueueBase* GetOutput(std::string_view name) override; + StageMetricProto FlushMetrics() override; + + Queue* output_queue() { return &primary_output_queue_; } + + private: + struct ThreadContext { + LoadMetricUpdater load_metric_updater; + }; + + void Worker(std::stop_token stop_token, ThreadContext* context); + + const ChunkUnpackerConfig config_; + const uint32_t run_seed_; + Queue primary_output_queue_; + std::optional> prefetch_output_queue_; + // thread_contexts_ must be declared before thread_pool_ to ensure + // thread_pool_ is destroyed first (stopping threads before contexts). + std::vector> thread_contexts_; + ThreadPool thread_pool_; +}; + +std::vector PickSampledPositions(int32_t n, double p, + int32_t iteration, + absl::BitGen& gen); + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/chunk_unpacker_test.cc b/csrc/loader/stages/chunk_unpacker_test.cc new file mode 100644 index 00000000..fe24bb16 --- /dev/null +++ b/csrc/loader/stages/chunk_unpacker_test.cc @@ -0,0 +1,269 @@ +#include "loader/stages/chunk_unpacker.h" + +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/container/flat_hash_set.h" +#include "absl/random/random.h" +#include "absl/random/seed_sequences.h" +#include "gtest/gtest.h" +#include "libs/lc0/src/trainingdata/trainingdata_v6.h" +#include "loader/stages/training_chunk.h" +#include "proto/data_loader_config.pb.h" +#include "utils/queue.h" + +namespace lczero { +namespace training { + +namespace { + +template +class PassthroughStage : public Stage { + public: + explicit PassthroughStage(Queue* queue) : queue_(queue) {} + + void Start() override {} + void Stop() override {} + StageMetricProto FlushMetrics() override { return StageMetricProto(); } + QueueBase* GetOutput(std::string_view name = "") override { + (void)name; + return queue_; + } + void SetInputs(absl::Span inputs) override { + if (!inputs.empty()) { + throw std::runtime_error("PassthroughStage expects no inputs"); + } + } + + private: + Queue* queue_; +}; + +} // namespace + +class ChunkUnpackerTest : public ::testing::Test { + protected: + void SetUp() override { + input_queue_ = std::make_unique>(10); + config_.set_threads(1); + config_.mutable_output()->set_queue_capacity(10); + config_.set_position_sampling_rate(1.0f); + } + + FrameType CreateTestFrame(uint32_t version) { + FrameType frame{}; + frame.version = version; + frame.input_format = 3; + frame.root_q = 0.5f; + return frame; + } + + TrainingChunk MakeChunk(std::vector frames, + std::string sort_key = "source", size_t index = 0, + uint32_t use = 0) { + TrainingChunk chunk; + chunk.sort_key = std::move(sort_key); + chunk.index_within_sort_key = index; + chunk.use_count = use; + chunk.frames = std::move(frames); + return chunk; + } + + std::unique_ptr> input_queue_; + ChunkUnpackerConfig config_; +}; + +TEST_F(ChunkUnpackerTest, UnpacksSingleFrame) { + ChunkUnpacker unpacker(config_); + unpacker.SetInputs({input_queue_.get()}); + unpacker.Start(); + + FrameType test_frame = CreateTestFrame(6); + auto producer = input_queue_->CreateProducer(); + producer.Put(MakeChunk({test_frame})); + producer.Close(); + + auto output_frame = unpacker.output_queue()->Get(); + EXPECT_EQ(output_frame.version, 6); + EXPECT_EQ(output_frame.input_format, 3); + EXPECT_EQ(output_frame.root_q, 0.5f); +} + +TEST_F(ChunkUnpackerTest, UnpacksMultipleFrames) { + ChunkUnpacker unpacker(config_); + unpacker.SetInputs({input_queue_.get()}); + unpacker.Start(); + + std::vector test_frames = {CreateTestFrame(6), CreateTestFrame(7), + CreateTestFrame(8)}; + auto producer = input_queue_->CreateProducer(); + producer.Put(MakeChunk(test_frames)); + producer.Close(); + + std::vector actual_versions; + actual_versions.reserve(test_frames.size()); + for (size_t i = 0; i < test_frames.size(); ++i) { + auto output_frame = unpacker.output_queue()->Get(); + actual_versions.push_back(output_frame.version); + EXPECT_EQ(output_frame.input_format, 3); + EXPECT_EQ(output_frame.root_q, 0.5f); + } + + std::vector expected_versions; + expected_versions.reserve(test_frames.size()); + for (const auto& frame : test_frames) { + expected_versions.push_back(frame.version); + } + + absl::c_sort(actual_versions); + absl::c_sort(expected_versions); + EXPECT_EQ(actual_versions, expected_versions); +} + +TEST_F(ChunkUnpackerTest, UnpacksMultipleChunks) { + ChunkUnpacker unpacker(config_); + unpacker.SetInputs({input_queue_.get()}); + unpacker.Start(); + + auto producer = input_queue_->CreateProducer(); + + // Send first chunk with 2 frames + std::vector chunk1_frames = {CreateTestFrame(10), + CreateTestFrame(11)}; + producer.Put(MakeChunk(chunk1_frames, "source", 0)); + + // Send second chunk with 1 frame + std::vector chunk2_frames = {CreateTestFrame(12)}; + producer.Put(MakeChunk(chunk2_frames, "source", 1)); + + producer.Close(); + + // Verify all frames are output + std::vector expected_versions = {10, 11, 12}; + std::vector actual_versions; + actual_versions.reserve(expected_versions.size()); + for (size_t i = 0; i < expected_versions.size(); ++i) { + auto output_frame = unpacker.output_queue()->Get(); + actual_versions.push_back(output_frame.version); + EXPECT_EQ(output_frame.input_format, 3); + EXPECT_EQ(output_frame.root_q, 0.5f); + } + + absl::c_sort(actual_versions); + absl::c_sort(expected_versions); + EXPECT_EQ(actual_versions, expected_versions); +} + +TEST_F(ChunkUnpackerTest, HandlesEmptyChunk) { + ChunkUnpacker unpacker(config_); + unpacker.SetInputs({input_queue_.get()}); + unpacker.Start(); + + auto producer = input_queue_->CreateProducer(); + TrainingChunk empty_chunk; + empty_chunk.sort_key = "source"; + empty_chunk.index_within_sort_key = 0; + producer.Put(std::move(empty_chunk)); + producer.Close(); + + // Should not produce any output frames, queue should close + EXPECT_THROW(unpacker.output_queue()->Get(), QueueClosedException); +} + +TEST_F(ChunkUnpackerTest, HandlesQueueClosure) { + ChunkUnpacker unpacker(config_); + unpacker.SetInputs({input_queue_.get()}); + unpacker.Start(); + + // Close input queue without sending data + input_queue_->Close(); + + // Output queue should eventually close + EXPECT_THROW(unpacker.output_queue()->Get(), QueueClosedException); +} + +TEST(PickSampledPositionsTest, Deterministic) { + absl::BitGen gen1(absl::SeedSeq{42}); + std::vector result1 = PickSampledPositions(1000, 0.1, 5, gen1); + + absl::BitGen gen2(absl::SeedSeq{42}); + std::vector result2 = PickSampledPositions(1000, 0.1, 5, gen2); + + EXPECT_EQ(result1, result2); +} + +TEST(PickSampledPositionsTest, FullBucketFirstRound) { + absl::BitGen gen(absl::SeedSeq{42}); + const uint32_t n = 10000; + const double p = 0.1; + std::vector result = PickSampledPositions(n, p, 0, gen); + // Expect size to be around n*p. + EXPECT_NEAR(result.size(), n * p, n * p * 0.25); +} + +TEST(PickSampledPositionsTest, DisjointBuckets) { + absl::BitGen gen(absl::SeedSeq{42}); + const uint32_t n = 1000; + const double p = 0.1; + + std::vector bucket1 = PickSampledPositions(n, p, 0, gen); + absl::c_sort(bucket1); + + // The generator state is now changed. For the next bucket, we need a fresh + // one with the same seed to test the logic for a different iteration. + absl::BitGen gen2(absl::SeedSeq{42}); + std::vector bucket2 = PickSampledPositions(n, p, 1, gen2); + absl::c_sort(bucket2); + + std::vector intersection; + absl::c_set_intersection(bucket1, bucket2, std::back_inserter(intersection)); + + EXPECT_TRUE(intersection.empty()); +} + +TEST(PickSampledPositionsTest, PartialBucketElementsAreReturned) { + absl::BitGen gen(absl::SeedSeq{42}); + const uint32_t n = 1000; + const double p = 0.8; // remainder 0.2 + + // In round 1, elements with toss >= 0.8 are for iteration 1. + absl::BitGen gen1(absl::SeedSeq{42}); + std::vector expected_from_round1; + for (uint32_t i = 0; i < n; ++i) { + double toss = absl::Uniform(gen1, 0.0, 1.0); + if (toss >= 0.8) { + expected_from_round1.push_back(i); + } + } + absl::c_sort(expected_from_round1); + + std::vector result = PickSampledPositions(n, p, 1, gen); + absl::c_sort(result); + + // Check if all elements from round 1 are in the final result. + // This will fail with the current implementation because they are discarded. + std::vector intersection; + absl::c_set_intersection(expected_from_round1, result, + std::back_inserter(intersection)); + + EXPECT_GT(expected_from_round1.size(), 50); // High probability for n=1000 + EXPECT_EQ(intersection.size(), expected_from_round1.size()); +} + +TEST(PickSampledPositionsTest, PartialBucketCompletedSize) { + absl::BitGen gen(absl::SeedSeq{42}); + const uint32_t n = 10000; + const double p = 0.8; // remainder 0.2 + + std::vector result = PickSampledPositions(n, p, 1, gen); + + // Expect size to be around n*p. + // This will fail due to incorrect probability calculation for completion. + EXPECT_NEAR(result.size(), n * p, n * p * 0.25); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/file_path_provider.cc b/csrc/loader/stages/file_path_provider.cc new file mode 100644 index 00000000..1897343d --- /dev/null +++ b/csrc/loader/stages/file_path_provider.cc @@ -0,0 +1,365 @@ +#include "loader/stages/file_path_provider.h" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "loader/data_loader_metrics.h" +#include "proto/data_loader_config.pb.h" + +namespace lczero { +namespace training { + +namespace { + +bool ShouldSkipName(std::string_view name) { + return !name.empty() && name.front() == '.'; +} + +bool ShouldSkipPathEntry(const FilePathProvider::Path& path) { + return ShouldSkipName(path.filename().string()); +} + +} // namespace + +FilePathProvider::FilePathProvider(const FilePathProviderConfig& config) + : SingleOutputStage(config.output()), + directory_(config.directory()), + producer_(output_queue()->CreateProducer()), + load_metric_updater_() { + LOG(INFO) << "Initializing FilePathProvider for directory: " + << config.directory(); + inotify_fd_ = inotify_init1(IN_CLOEXEC | IN_NONBLOCK); + CHECK_NE(inotify_fd_, -1) + << "Failed to initialize inotify: " << strerror(errno); +} + +FilePathProvider::~FilePathProvider() { + LOG(INFO) << "FilePathProvider shutting down."; + Stop(); + if (inotify_fd_ != -1) close(inotify_fd_); + LOG(INFO) << "FilePathProvider shutdown complete."; +} + +void FilePathProvider::SetInputs(absl::Span inputs) { + if (!inputs.empty()) { + throw std::runtime_error( + "FilePathProvider expects no inputs, but received " + + std::to_string(inputs.size())); + } +} + +void FilePathProvider::Start() { + LOG(INFO) << "Starting FilePathProvider monitoring thread."; + thread_pool_.Enqueue( + [this](std::stop_token stop_token) { Worker(stop_token); }); +} + +void FilePathProvider::Stop() { + if (stop_source_.stop_requested()) return; + LOG(INFO) << "Stopping FilePathProvider."; + LOG(INFO) << "Stopping all watches..."; + for (const auto& [wd, path] : watch_descriptors_) { + inotify_rm_watch(inotify_fd_, wd); + } + watch_descriptors_.clear(); + stop_source_.request_stop(); + thread_pool_.Shutdown(); + producer_.Close(); +} + +StageMetricProto FilePathProvider::FlushMetrics() { + StageMetricProto stage_metric; + auto load_metrics = load_metric_updater_.FlushMetrics(); + load_metrics.set_name("load"); + *stage_metric.add_load_metrics() = std::move(load_metrics); + *stage_metric.add_queue_metrics() = + MetricsFromQueue("output", *output_queue()); + return stage_metric; +} + +void FilePathProvider::AddDirectory(const Path& directory, + std::stop_token stop_token) { + ScanDirectoryWithWatch(directory, stop_token); + + LOG(INFO) << "FilePathProvider registered " << directory + << "; active watch descriptors: " << watch_descriptors_.size(); + + // Signal that initial scan is complete + LOG(INFO) << "FilePathProvider initial scan complete"; + producer_.Put( + {{.filepath = Path{}, .message_type = MessageType::kInitialScanComplete}}, + stop_token); +} + +void FilePathProvider::ScanDirectoryWithWatch(const Path& directory, + std::stop_token stop_token) { + // Step 1: Set up watch first + int wd = inotify_add_watch(inotify_fd_, directory.c_str(), + IN_CLOSE_WRITE | IN_MOVED_TO | IN_CREATE | + IN_DELETE | IN_DELETE_SELF | IN_MOVE); + CHECK_NE(wd, -1) << "Failed to add inotify watch for " << directory << ": " + << strerror(errno); + watch_descriptors_[wd] = directory; + + // Step 2: Scan directory non-recursively, remembering files and subdirs + std::vector files; + std::vector subdirectories; + std::error_code ec; + auto iterator = std::filesystem::directory_iterator(directory, ec); + CHECK(!ec) << "Failed to iterate directory " << directory << ": " + << ec.message(); + + for (const auto& entry : iterator) { + const Path entry_path = entry.path(); + if (ShouldSkipPathEntry(entry_path)) continue; + + if (entry.is_regular_file(ec) && !ec) { + files.push_back(entry_path); + } else if (entry.is_directory(ec) && !ec) { + subdirectories.push_back(entry_path); + } + } + + const size_t initial_file_count = files.size(); + const size_t subdirectory_count = subdirectories.size(); + LOG(INFO) << "FilePathProvider scanned " << directory << " discovering " + << initial_file_count << " file(s) and " << subdirectory_count + << " subdirectory(ies) before watch reconciliation."; + + // Send notifications for discovered files + constexpr size_t kBatchSize = 10000; + std::vector batch; + batch.reserve(kBatchSize); + + auto flush_batch = [&]() { + if (batch.empty()) return; + producer_.Put(batch, stop_token); + batch.clear(); + }; + + for (const auto& filepath : files) { + batch.push_back( + {.filepath = filepath.string(), .message_type = MessageType::kFile}); + if (batch.size() >= kBatchSize) flush_batch(); + } + + if (initial_file_count > 0) { + LOG(INFO) << "FilePathProvider enqueued " << initial_file_count + << " file(s) from initial scan of " << directory; + } + + // Step 3: Read from watch descriptor, skipping already discovered files + ProcessWatchEventsForNewItems(files); + + // Step 4: Clean the files vector to save memory + files.clear(); + + // Step 5: Recursively call for subdirectories + for (const auto& subdir : subdirectories) { + if (stop_token.stop_requested()) return; + ScanDirectoryWithWatch(subdir, stop_token); + } + + // Flush any remaining files + flush_batch(); +} + +void FilePathProvider::ProcessWatchEventsForNewItems( + const std::vector& known_files) { + // Create a set for fast lookup of already discovered files + absl::flat_hash_set known_file_set; + for (const auto& file : known_files) { + known_file_set.insert(file.string()); + } + + // Process any events that may have occurred during scanning + std::array buffer; + std::vector new_files; + + while (true) { + ssize_t length = read(inotify_fd_, buffer.data(), buffer.size()); + if (length <= 0) break; // No more events to process + + ssize_t offset = 0; + while (offset < length) { + const struct inotify_event* event = + reinterpret_cast(buffer.data() + offset); + + const bool skip_entry = event->len > 0 && ShouldSkipName(event->name); + + // Only process file creation/write events, skip already known files + if ((event->mask & (IN_CLOSE_WRITE | IN_MOVED_TO)) != 0 && + event->len > 0 && !skip_entry) { + const Path directory(watch_descriptors_.at(event->wd)); + Path filepath = directory / event->name; + std::string filepath_string = filepath.string(); + + // Only add if we haven't seen this file before + if (!known_file_set.contains(filepath_string)) { + new_files.push_back({.filepath = std::move(filepath_string), + .message_type = MessageType::kFile}); + } + } + + offset += sizeof(struct inotify_event) + event->len; + } + } + + // Send notifications for any new files discovered through watch events + if (!new_files.empty()) { + LOG(INFO) << "FilePathProvider observed " << new_files.size() + << " new file(s) while reconciling race events."; + producer_.Put(new_files); + } +} + +void FilePathProvider::AddWatchRecursive(const Path& path) { + // Add watch for current directory + int wd = inotify_add_watch(inotify_fd_, path.c_str(), + IN_CLOSE_WRITE | IN_MOVED_TO | IN_CREATE | + IN_DELETE | IN_DELETE_SELF | IN_MOVE); + CHECK_NE(wd, -1) << "Failed to add inotify watch for " << path << ": " + << strerror(errno); + watch_descriptors_[wd] = path; + + // Recursively add watches for subdirectories + std::error_code ec; + auto iterator = std::filesystem::directory_iterator(path, ec); + CHECK(!ec) << "Failed to iterate directory " << path << ": " << ec.message(); + + for (const auto& entry : iterator) { + const Path entry_path = entry.path(); + if (ShouldSkipPathEntry(entry_path)) continue; + if (!entry.is_directory(ec) || ec) continue; + AddWatchRecursive(entry_path); + } +} + +void FilePathProvider::RemoveWatchRecursive(const Path& base) { + absl::erase_if(watch_descriptors_, [&](const auto& pair) { + const auto& [wd, path] = pair; + const auto mismatch_iter = absl::c_mismatch(base, path).first; + // If path is not a subdirectory (or equal) of base, skip. + if (mismatch_iter != base.end()) return false; + inotify_rm_watch(inotify_fd_, wd); + return true; + }); +} + +void FilePathProvider::Worker(std::stop_token stop_token) { + // Perform directory scanning in background thread + AddDirectory(directory_, stop_token); + + int epoll_fd = epoll_create1(EPOLL_CLOEXEC); + CHECK_NE(epoll_fd, -1) << "Failed to create epoll fd: " << strerror(errno); + absl::Cleanup epoll_cleanup([epoll_fd]() { close(epoll_fd); }); + + struct epoll_event event; + event.events = EPOLLIN; + event.data.fd = inotify_fd_; + CHECK_EQ(epoll_ctl(epoll_fd, EPOLL_CTL_ADD, inotify_fd_, &event), 0) + << "Failed to add inotify fd to epoll: " << strerror(errno); + + while (!stop_token.stop_requested()) { + { + LoadMetricPauser pauser(load_metric_updater_); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + if (stop_token.stop_requested()) { + pauser.DoNotResume(); + break; + } + } + + struct epoll_event event; + int nfds = epoll_wait(epoll_fd, &event, 1, 0); // Non-blocking check + CHECK_NE(nfds, -1) << "epoll_wait failed: " << strerror(errno); + if (nfds == 0) continue; // No events. + + do { + assert(nfds == 1 && event.data.fd == inotify_fd_); + ProcessInotifyEvents(producer_, stop_token); + nfds = epoll_wait(epoll_fd, &event, 1, 0); + } while (nfds > 0); + } +} + +void FilePathProvider::ProcessInotifyEvents(Queue::Producer& producer, + std::stop_token stop_token) { + constexpr size_t kNotifyBatchSize = 10000; + std::vector files; + std::array buffer; + + auto flush_batch = [&]() { + if (files.empty()) return; + producer.Put(files, stop_token); + files.clear(); + }; + + while (true) { + ssize_t length = read(inotify_fd_, buffer.data(), buffer.size()); + if (length <= 0) break; // No more events to process + + ssize_t offset = 0; + while (offset < length) { + const struct inotify_event* event = + reinterpret_cast(buffer.data() + offset); + auto file = ProcessInotifyEvent(*event, stop_token); + if (file) files.push_back(*file); + if (files.size() >= kNotifyBatchSize) flush_batch(); + offset += sizeof(struct inotify_event) + event->len; + } + } + + flush_batch(); // Flush any remaining files in the batch +} + +auto FilePathProvider::ProcessInotifyEvent(const struct inotify_event& event, + std::stop_token stop_token) + -> std::optional { + if (event.mask & IN_IGNORED) return std::nullopt; + + const Path directory(watch_descriptors_.at(event.wd)); + const bool has_name = event.len > 0 && event.name[0] != '\0'; + const bool skip_entry = has_name && ShouldSkipName(event.name); + const Path filepath = has_name ? directory / event.name : directory; + + // Handle different event types + if ((event.mask & (IN_CLOSE_WRITE | IN_MOVED_TO)) != 0 && has_name && + !skip_entry) { + // File finished writing or moved into directory + return File{.filepath = filepath, .message_type = MessageType::kFile}; + } + + constexpr uint32_t kDirCreateMask = IN_CREATE | IN_ISDIR; + constexpr uint32_t kDirDeleteMask = IN_DELETE | IN_ISDIR; + if ((event.mask & kDirCreateMask) == kDirCreateMask) { + if (!has_name || skip_entry) return std::nullopt; + ScanDirectoryWithWatch(filepath, stop_token); + } else if ((event.mask & kDirDeleteMask) == kDirDeleteMask) { + if (!has_name || skip_entry) return std::nullopt; + // Directory deleted - remove all watches for it and subdirectories + RemoveWatchRecursive(filepath); + } else if (event.mask & IN_DELETE_SELF) { + RemoveWatchRecursive(directory); + } + + return std::nullopt; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/file_path_provider.h b/csrc/loader/stages/file_path_provider.h new file mode 100644 index 00000000..bc7dadd6 --- /dev/null +++ b/csrc/loader/stages/file_path_provider.h @@ -0,0 +1,93 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "loader/stages/stage.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/metrics/load_metric.h" +#include "utils/metrics/printer.h" +#include "utils/metrics/statistics_metric.h" +#include "utils/queue.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +// Message types for FilePathProvider output. +enum class FilePathProviderMessageType { + kFile, // File discovered (initial scan or inotify) + kInitialScanComplete // Initial scan is complete (empty filepath) +}; + +// Output type for FilePathProvider. +struct FilePathProviderFile { + std::filesystem::path filepath; + FilePathProviderMessageType message_type; +}; + +// This class watches for new files in a directory (recursively) and notifies +// registered observers when new files are either closed after writing or +// renamed into. +// Uses background thread to monitor the directory. +class FilePathProvider : public SingleOutputStage { + public: + using Path = std::filesystem::path; + using MessageType = FilePathProviderMessageType; + using File = FilePathProviderFile; + + explicit FilePathProvider(const FilePathProviderConfig& config); + ~FilePathProvider(); + + // Starts monitoring the directory + void Start() override; + + // Closes the output queue, signaling completion + void Stop() override; + + // Returns current metrics and clears them. + StageMetricProto FlushMetrics() override; + + // FilePathProvider has no inputs. + void SetInputs(absl::Span inputs) override; + + private: + // Starts monitoring the directory. + void AddDirectory(const Path& directory, std::stop_token stop_token); + + void Worker(std::stop_token stop_token); + void AddWatchRecursive(const Path& path); + void RemoveWatchRecursive(const Path& path); + void ScanDirectoryWithWatch(const Path& directory, + std::stop_token stop_token); + void ProcessWatchEventsForNewItems(const std::vector& known_files); + void ProcessInotifyEvents(Queue::Producer& producer, + std::stop_token stop_token); + std::optional ProcessInotifyEvent(const struct inotify_event& event, + std::stop_token stop_token); + + int inotify_fd_; + // Watch descriptor to directory path. + absl::flat_hash_map watch_descriptors_; + + Path directory_; // Directory to monitor + Queue::Producer producer_; + + LoadMetricUpdater load_metric_updater_; + std::stop_source stop_source_; + ThreadPool thread_pool_{1, ThreadPoolOptions{}, stop_source_}; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/file_path_provider_main.cc b/csrc/loader/stages/file_path_provider_main.cc new file mode 100644 index 00000000..5505c3ab --- /dev/null +++ b/csrc/loader/stages/file_path_provider_main.cc @@ -0,0 +1,59 @@ +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "loader/stages/file_path_provider.h" + +ABSL_FLAG(std::string, directory, "", "Directory to monitor for files"); + +int main(int argc, char* argv[]) { + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo); + + std::string directory = absl::GetFlag(FLAGS_directory); + if (directory.empty()) { + std::cerr << "Usage: " << argv[0] << " --directory=" + << std::endl; + return 1; + } + + LOG(INFO) << "Starting to monitor directory: " << directory; + lczero::training::FilePathProviderConfig config; + config.mutable_output()->set_queue_capacity(16); + config.set_directory(directory); + lczero::training::FilePathProvider file_path_provider(config); + + // Consumer thread to read from the queue + std::thread consumer_thread([&file_path_provider]() { + auto* queue = file_path_provider.output_queue(); + try { + while (true) { + auto file = queue->Get(); + const char* type_str = + (file.message_type == + lczero::training::FilePathProvider::MessageType::kFile) + ? "File" + : "Initial scan complete"; + LOG(INFO) << "File " << type_str << ": " << file.filepath; + } + } catch (const lczero::QueueClosedException&) { + LOG(INFO) << "Queue closed, consumer thread exiting"; + } + }); + + LOG(INFO) << "Monitoring for files... Press Enter to exit."; + std::cin.get(); + + // Close the queue and wait for consumer to finish + file_path_provider.Stop(); + consumer_thread.join(); + + return 0; +} diff --git a/csrc/loader/stages/file_path_provider_test.cc b/csrc/loader/stages/file_path_provider_test.cc new file mode 100644 index 00000000..81461c64 --- /dev/null +++ b/csrc/loader/stages/file_path_provider_test.cc @@ -0,0 +1,235 @@ +#include "loader/stages/file_path_provider.h" + +#include + +#include +#include +#include +#include +#include +#include + +namespace lczero { +namespace training { + +namespace { + +FilePathProviderConfig MakeConfig(const std::filesystem::path& directory) { + FilePathProviderConfig config; + config.mutable_output()->set_queue_capacity(128); + config.set_directory(directory.string()); + return config; +} + +std::string RelativeTo(const std::filesystem::path& base, + const std::filesystem::path& target) { + return target.lexically_relative(base).generic_string(); +} + +} // namespace + +class FilePathProviderTest : public ::testing::Test { + protected: + void SetUp() override { + test_dir_ = + std::filesystem::temp_directory_path() / + ("file_path_provider_test_" + + std::to_string( + std::chrono::steady_clock::now().time_since_epoch().count())); + std::filesystem::create_directories(test_dir_); + } + + void TearDown() override { + if (std::filesystem::exists(test_dir_)) { + std::filesystem::remove_all(test_dir_); + } + } + + void CreateFile(const std::filesystem::path& path, + const std::string& content = "payload") { + std::filesystem::create_directories(path.parent_path()); + std::ofstream file(path); + file << content; + } + + void CreateDirectory(const std::filesystem::path& path) { + std::filesystem::create_directories(path); + } + + std::vector DrainInitialScan( + Queue* queue) { + std::vector files; + while (true) { + auto message = queue->Get(); + if (message.message_type == + FilePathProvider::MessageType::kInitialScanComplete) { + EXPECT_TRUE(message.filepath.empty()); + break; + } + if (message.message_type != FilePathProvider::MessageType::kFile) { + ADD_FAILURE() << "Unexpected message type in initial scan."; + continue; + } + files.push_back(message.filepath); + } + return files; + } + + FilePathProvider::File AwaitNextFile(Queue* queue) { + while (true) { + auto message = queue->Get(); + if (message.message_type == FilePathProvider::MessageType::kFile) { + return message; + } + if (message.message_type != + FilePathProvider::MessageType::kInitialScanComplete) { + ADD_FAILURE() + << "Unexpected message type while waiting for file notification."; + } + } + } + + FilePathProviderConfig Config() const { return MakeConfig(test_dir_); } + + std::filesystem::path test_dir_; +}; + +TEST_F(FilePathProviderTest, ConstructorCreatesQueue) { + FilePathProvider provider(Config()); + provider.Start(); + + auto* queue = provider.output_queue(); + ASSERT_NE(queue, nullptr); + EXPECT_EQ(queue->Capacity(), 128); + + auto message = queue->Get(); + EXPECT_EQ(message.message_type, + FilePathProvider::MessageType::kInitialScanComplete); + EXPECT_TRUE(message.filepath.empty()); + + provider.Stop(); +} + +TEST_F(FilePathProviderTest, InitialScanFindsVisibleFiles) { + CreateFile(test_dir_ / "file1.txt"); + CreateFile(test_dir_ / "file2.txt"); + CreateFile(test_dir_ / "sub" / "nested.txt"); + + FilePathProvider provider(Config()); + provider.Start(); + auto* queue = provider.output_queue(); + + auto discovered = DrainInitialScan(queue); + std::unordered_set relative_paths; + for (const auto& path : discovered) { + relative_paths.insert(RelativeTo(test_dir_, path)); + } + + EXPECT_EQ(relative_paths.size(), 3u); + EXPECT_TRUE(relative_paths.count("file1.txt")); + EXPECT_TRUE(relative_paths.count("file2.txt")); + EXPECT_TRUE(relative_paths.count("sub/nested.txt")); + + provider.Stop(); +} + +TEST_F(FilePathProviderTest, InitialScanSkipsHiddenEntries) { + CreateFile(test_dir_ / "visible.txt"); + CreateFile(test_dir_ / ".hidden_file"); + CreateFile(test_dir_ / ".hidden_dir" / "nested.txt"); + CreateFile(test_dir_ / "visible_dir" / "child.txt"); + + FilePathProvider provider(Config()); + provider.Start(); + auto* queue = provider.output_queue(); + + auto discovered = DrainInitialScan(queue); + std::unordered_set relative_paths; + for (const auto& path : discovered) { + relative_paths.insert(RelativeTo(test_dir_, path)); + } + + EXPECT_TRUE(relative_paths.count("visible.txt")); + EXPECT_TRUE(relative_paths.count("visible_dir/child.txt")); + EXPECT_FALSE(relative_paths.count(".hidden_file")); + EXPECT_FALSE(relative_paths.count(".hidden_dir/nested.txt")); + + provider.Stop(); +} + +TEST_F(FilePathProviderTest, DetectsNewVisibleFile) { + FilePathProvider provider(Config()); + provider.Start(); + auto* queue = provider.output_queue(); + DrainInitialScan(queue); + + CreateFile(test_dir_ / "new_file.txt"); + + auto message = AwaitNextFile(queue); + EXPECT_EQ(RelativeTo(test_dir_, message.filepath), "new_file.txt"); + + provider.Stop(); +} + +TEST_F(FilePathProviderTest, DetectsFilesInPreExistingSubdirectory) { + auto subdir = test_dir_ / "subdir"; + CreateDirectory(subdir); + + FilePathProvider provider(Config()); + provider.Start(); + auto* queue = provider.output_queue(); + DrainInitialScan(queue); + + CreateFile(subdir / "from_subdir.txt"); + + auto message = AwaitNextFile(queue); + EXPECT_EQ(RelativeTo(test_dir_, message.filepath), "subdir/from_subdir.txt"); + + provider.Stop(); +} + +TEST_F(FilePathProviderTest, IgnoresHiddenFileEvents) { + FilePathProvider provider(Config()); + provider.Start(); + auto* queue = provider.output_queue(); + DrainInitialScan(queue); + + CreateFile(test_dir_ / ".hidden_event.txt"); + CreateFile(test_dir_ / "visible_after_hidden.txt"); + + auto message = AwaitNextFile(queue); + EXPECT_EQ(RelativeTo(test_dir_, message.filepath), + "visible_after_hidden.txt"); + + provider.Stop(); +} + +TEST_F(FilePathProviderTest, SkipsHiddenDirectoryRecursion) { + FilePathProvider provider(Config()); + provider.Start(); + auto* queue = provider.output_queue(); + DrainInitialScan(queue); + + CreateDirectory(test_dir_ / ".hidden_dir"); + CreateFile(test_dir_ / ".hidden_dir" / "inner.txt"); + CreateFile(test_dir_ / "outer.txt"); + + auto message = AwaitNextFile(queue); + EXPECT_EQ(RelativeTo(test_dir_, message.filepath), "outer.txt"); + + provider.Stop(); +} + +TEST_F(FilePathProviderTest, HandlesEmptyDirectory) { + FilePathProvider provider(Config()); + provider.Start(); + auto* queue = provider.output_queue(); + + auto discovered = DrainInitialScan(queue); + EXPECT_TRUE(discovered.empty()); + + provider.Stop(); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/join_stage.cc b/csrc/loader/stages/join_stage.cc new file mode 100644 index 00000000..d1524681 --- /dev/null +++ b/csrc/loader/stages/join_stage.cc @@ -0,0 +1,89 @@ +#include "loader/stages/join_stage.h" + +#include + +#include "loader/data_loader_metrics.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/metrics/load_metric.h" + +namespace lczero { +namespace training { + +template +JoinStage::JoinStage(const JoinPositionsConfig& config) + : SingleOutputStage(config.output()) {} + +template +JoinStage::~JoinStage() { + Stop(); +} + +template +void JoinStage::SetInputs(absl::Span inputs) { + input_queues_.clear(); + for (QueueBase* base_queue : inputs) { + auto* typed_queue = dynamic_cast*>(base_queue); + if (!typed_queue) throw std::runtime_error("Input queue type mismatch"); + input_queues_.push_back(typed_queue); + } +} + +template +void JoinStage::Start() { + thread_contexts_.clear(); + thread_pool_ = std::make_unique(input_queues_.size()); + for (size_t i = 0; i < input_queues_.size(); ++i) { + thread_contexts_.push_back(std::make_unique()); + } + for (size_t i = 0; i < input_queues_.size(); ++i) { + thread_pool_->Enqueue([this, i](std::stop_token stop_token) { + Worker(stop_token, input_queues_[i], thread_contexts_[i].get()); + }); + } +} + +template +void JoinStage::Worker(std::stop_token stop_token, Queue* input_queue, + ThreadContext* context) { + auto producer = this->output_queue()->CreateProducer(); + try { + while (true) { + auto item = [&]() { + LoadMetricPauser pauser(context->load_metric_updater); + return input_queue->Get(stop_token); + }(); + producer.Put(std::move(item), stop_token); + } + } catch (const QueueClosedException&) { + } +} + +template +void JoinStage::Stop() { + if (!thread_pool_ || thread_pool_->stop_token().stop_requested()) return; + LOG(INFO) << "Stopping JoinStage."; + thread_pool_->Shutdown(); + this->output_queue()->Close(); +} + +template +StageMetricProto JoinStage::FlushMetrics() { + StageMetricProto metrics; + LoadMetricProto aggregated_load; + aggregated_load.set_name("load"); + for (const auto& context : thread_contexts_) { + UpdateFrom(aggregated_load, context->load_metric_updater.FlushMetrics()); + } + *metrics.add_load_metrics() = std::move(aggregated_load); + *metrics.add_queue_metrics() = + MetricsFromQueue("output", *this->output_queue()); + + return metrics; +} + +// Explicit template instantiation for FrameType. +template class JoinStage; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/join_stage.h b/csrc/loader/stages/join_stage.h new file mode 100644 index 00000000..8cedbc5a --- /dev/null +++ b/csrc/loader/stages/join_stage.h @@ -0,0 +1,54 @@ +// ABOUTME: Stage that joins multiple input queues into a single output. +// ABOUTME: Spawns one thread per input to read and forward items. +#pragma once + +#include +#include +#include +#include + +#include "loader/data_loader_metrics.h" +#include "loader/frame_type.h" +#include "loader/stages/stage.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/queue.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +// Template stage that joins multiple input queues into a single output. +// Spawns one thread per input to consume and forward items. +template +class JoinStage : public SingleOutputStage { + public: + using OutputType = T; + + explicit JoinStage(const JoinPositionsConfig& config); + ~JoinStage(); + + void Start() override; + void Stop() override; + StageMetricProto FlushMetrics() override; + void SetInputs(absl::Span inputs) override; + + private: + struct ThreadContext { + LoadMetricUpdater load_metric_updater; + }; + + void Worker(std::stop_token stop_token, Queue* input_queue, + ThreadContext* context); + + std::vector*> input_queues_; + // thread_contexts_ must be declared before thread_pool_ to ensure + // thread_pool_ is destroyed first (stopping threads before contexts). + std::vector> thread_contexts_; + std::unique_ptr thread_pool_; +}; + +using JoinPositions = JoinStage; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/join_stage_test.cc b/csrc/loader/stages/join_stage_test.cc new file mode 100644 index 00000000..9e9cbcc2 --- /dev/null +++ b/csrc/loader/stages/join_stage_test.cc @@ -0,0 +1,140 @@ +#include "loader/stages/join_stage.h" + +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "gtest/gtest.h" +#include "libs/lc0/src/trainingdata/trainingdata_v6.h" +#include "proto/data_loader_config.pb.h" +#include "utils/queue.h" + +namespace lczero { +namespace training { + +class JoinStageTest : public ::testing::Test { + protected: + void SetUp() override { config_.mutable_output()->set_queue_capacity(100); } + + FrameType CreateTestFrame(uint32_t version) { + FrameType frame{}; + frame.version = version; + frame.input_format = 3; + frame.root_q = 0.5f; + return frame; + } + + JoinPositionsConfig config_; +}; + +TEST_F(JoinStageTest, JoinsTwoInputs) { + auto input_queue_1 = std::make_unique>(10); + auto input_queue_2 = std::make_unique>(10); + + JoinPositions join_stage(config_); + join_stage.SetInputs({input_queue_1.get(), input_queue_2.get()}); + join_stage.Start(); + + auto producer_1 = input_queue_1->CreateProducer(); + auto producer_2 = input_queue_2->CreateProducer(); + + producer_1.Put(CreateTestFrame(1)); + producer_1.Put(CreateTestFrame(2)); + producer_2.Put(CreateTestFrame(3)); + producer_2.Put(CreateTestFrame(4)); + + absl::flat_hash_set received_versions; + for (int i = 0; i < 4; ++i) { + auto frame = join_stage.output_queue()->Get(); + received_versions.insert(frame.version); + } + + producer_1.Close(); + producer_2.Close(); + join_stage.Stop(); + + EXPECT_EQ(received_versions.size(), 4u); + EXPECT_TRUE(received_versions.contains(1)); + EXPECT_TRUE(received_versions.contains(2)); + EXPECT_TRUE(received_versions.contains(3)); + EXPECT_TRUE(received_versions.contains(4)); +} + +TEST_F(JoinStageTest, JoinsThreeInputs) { + auto input_queue_1 = std::make_unique>(10); + auto input_queue_2 = std::make_unique>(10); + auto input_queue_3 = std::make_unique>(10); + + JoinPositions join_stage(config_); + join_stage.SetInputs( + {input_queue_1.get(), input_queue_2.get(), input_queue_3.get()}); + join_stage.Start(); + + auto producer_1 = input_queue_1->CreateProducer(); + auto producer_2 = input_queue_2->CreateProducer(); + auto producer_3 = input_queue_3->CreateProducer(); + + producer_1.Put(CreateTestFrame(10)); + producer_2.Put(CreateTestFrame(20)); + producer_3.Put(CreateTestFrame(30)); + + absl::flat_hash_set received_versions; + for (int i = 0; i < 3; ++i) { + auto frame = join_stage.output_queue()->Get(); + received_versions.insert(frame.version); + } + + producer_1.Close(); + producer_2.Close(); + producer_3.Close(); + join_stage.Stop(); + + EXPECT_EQ(received_versions.size(), 3u); + EXPECT_TRUE(received_versions.contains(10)); + EXPECT_TRUE(received_versions.contains(20)); + EXPECT_TRUE(received_versions.contains(30)); +} + +TEST_F(JoinStageTest, HandlesEmptyInputs) { + auto input_queue_1 = std::make_unique>(10); + auto input_queue_2 = std::make_unique>(10); + + JoinPositions join_stage(config_); + join_stage.SetInputs({input_queue_1.get(), input_queue_2.get()}); + join_stage.Start(); + + auto producer_1 = input_queue_1->CreateProducer(); + auto producer_2 = input_queue_2->CreateProducer(); + + producer_1.Close(); + producer_2.Close(); + + auto maybe_frame = join_stage.output_queue()->MaybeGet(); + EXPECT_FALSE(maybe_frame.has_value()); + + join_stage.Stop(); +} + +TEST_F(JoinStageTest, FlushesMetrics) { + auto input_queue = std::make_unique>(10); + + JoinPositions join_stage(config_); + join_stage.SetInputs({input_queue.get()}); + join_stage.Start(); + + auto producer = input_queue->CreateProducer(); + producer.Put(CreateTestFrame(1)); + + auto frame = join_stage.output_queue()->Get(); + EXPECT_EQ(frame.version, 1u); + + producer.Close(); + join_stage.Stop(); + + auto metrics = join_stage.FlushMetrics(); + EXPECT_EQ(metrics.load_metrics_size(), 1); + EXPECT_EQ(metrics.queue_metrics_size(), 1); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/position_sampling.cc b/csrc/loader/stages/position_sampling.cc new file mode 100644 index 00000000..c8f3ec9e --- /dev/null +++ b/csrc/loader/stages/position_sampling.cc @@ -0,0 +1,26 @@ +#include "loader/stages/position_sampling.h" + +#include + +namespace lczero { +namespace training { + +float ComputePositionSamplingWeight(const FrameType& frame, + const PositionSamplingConfig& config) { + if (!config.has_diff_focus_q_weight() && !config.has_diff_focus_pol_scale()) { + return config.default_weight(); + } + if (std::isnan(frame.orig_q)) return config.default_weight(); + const float diff_q = std::abs(frame.best_q - frame.orig_q); + const float q_weight = config.diff_focus_q_weight(); + const float pol_scale = config.diff_focus_pol_scale(); + const float total = + (q_weight * diff_q + frame.policy_kld) / (q_weight + pol_scale); + return std::min( + std::pow(total * config.diff_focus_alpha() + config.diff_focus_beta(), + config.diff_focus_gamma()), + config.diff_focus_tau()); +} + +} // namespace training +} // namespace lczero \ No newline at end of file diff --git a/csrc/loader/stages/position_sampling.h b/csrc/loader/stages/position_sampling.h new file mode 100644 index 00000000..a5243550 --- /dev/null +++ b/csrc/loader/stages/position_sampling.h @@ -0,0 +1,13 @@ +#pragma once + +#include "loader/frame_type.h" +#include "proto/data_loader_config.pb.h" + +namespace lczero { +namespace training { + +float ComputePositionSamplingWeight(const FrameType& frame, + const PositionSamplingConfig& config); + +} // namespace training +} // namespace lczero \ No newline at end of file diff --git a/csrc/loader/stages/shuffling_chunk_pool.cc b/csrc/loader/stages/shuffling_chunk_pool.cc new file mode 100644 index 00000000..65e3051a --- /dev/null +++ b/csrc/loader/stages/shuffling_chunk_pool.cc @@ -0,0 +1,917 @@ +#include "loader/stages/shuffling_chunk_pool.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "loader/chunk_source/chunk_source.h" +#include "loader/data_loader_metrics.h" +#include "loader/stages/chunk_source_loader.h" +#include "loader/stages/position_sampling.h" +#include "proto/data_loader_config.pb.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +thread_local absl::BitGen ShufflingChunkPool::bitgen_{absl::MakeSeedSeq()}; + +ShufflingChunkPool::ShufflingChunkPool(const ShufflingChunkPoolConfig& config) + : primary_output_name_(config.output().name()), + primary_output_queue_( + config.output().queue_capacity(), + ToOverflowBehavior(config.output().overflow_behavior())), + chunk_pool_size_(config.chunk_pool_size()), + config_(config), + source_ingestion_pool_(config.source_ingestion_threads(), + ThreadPoolOptions{}, stop_source_), + chunk_loading_pool_(config.chunk_loading_threads(), ThreadPoolOptions{}, + stop_source_), + caching_pool_(config.has_cachehit_output() ? config.caching_threads() : 0, + ThreadPoolOptions{}, stop_source_) { + if (config.has_cachehit_output()) { + cachehit_output_name_ = config.cachehit_output().name(); + cachehit_output_queue_.emplace( + config.cachehit_output().queue_capacity(), + ToOverflowBehavior(config.cachehit_output().overflow_behavior())); + if (primary_output_name_ == *cachehit_output_name_) { + throw std::runtime_error(absl::StrCat( + "ShufflingChunkPool output names must be different, got: '", + primary_output_name_, "'")); + } + } + LOG(INFO) << "Initializing ShufflingChunkPool with pool size " + << config.chunk_pool_size(); +} + +ShufflingChunkPool::~ShufflingChunkPool() { Stop(); } + +void ShufflingChunkPool::SetInputs(absl::Span inputs) { + if (inputs.size() != 1 && inputs.size() != 2) { + throw std::runtime_error(absl::StrCat( + "ShufflingChunkPool expects 1 or 2 inputs, got ", inputs.size())); + } + if (inputs.size() == 2 && !cachehit_output_queue_.has_value()) { + throw std::runtime_error( + "ShufflingChunkPool received 2 inputs but cachehit_output is not " + "configured"); + } + if (inputs.size() == 1 && cachehit_output_queue_.has_value()) { + throw std::runtime_error( + "ShufflingChunkPool has cachehit_output configured but received only " + "1 input"); + } + primary_input_queue_ = dynamic_cast*>(inputs[0]); + if (!primary_input_queue_) { + throw std::runtime_error("ShufflingChunkPool primary input type mismatch"); + } + if (inputs.size() == 2) { + cache_request_queue_ = dynamic_cast*>(inputs[1]); + if (!cache_request_queue_) { + throw std::runtime_error( + "ShufflingChunkPool cache request input type mismatch"); + } + } +} + +QueueBase* ShufflingChunkPool::GetOutput(std::string_view name) { + if (name == primary_output_name_) return &primary_output_queue_; + if (cachehit_output_name_.has_value() && name == *cachehit_output_name_) { + return &*cachehit_output_queue_; + } + std::string available = absl::StrCat("'", primary_output_name_, "'"); + if (cachehit_output_name_.has_value()) { + absl::StrAppend(&available, ", '", *cachehit_output_name_, "'"); + } + throw std::runtime_error(absl::StrCat("ShufflingChunkPool unknown output '", + name, + "'. Available outputs: ", available)); +} + +void ShufflingChunkPool::Start() { + LOG(INFO) << "Starting ShufflingChunkPool initialization thread."; + initialization_thread_ = std::jthread([this]() { + try { + LOG(INFO) << "Starting ShufflingChunkPool with pool size " + << config_.chunk_pool_size(); + std::vector> uninitialized_sources = + InitializeChunkSources(); + ProcessInputFiles(std::move(uninitialized_sources)); + + // Start input processing worker that continuously processes new files. + for (size_t i = 0; i < source_ingestion_pool_.num_threads(); ++i) { + auto* context = + source_ingestion_thread_contexts_ + .emplace_back(std::make_unique()) + .get(); + source_ingestion_pool_.Enqueue( + [this, context](std::stop_token stop_token) { + SourceIngestionWorker(stop_token, context); + }); + } + + // Start output workers after everything is fully initialized. + LOG(INFO) << "ShufflingChunkPool initialization done, starting workers"; + for (size_t i = 0; i < chunk_loading_pool_.num_threads(); ++i) { + auto* context = + chunk_loading_thread_contexts_ + .emplace_back(std::make_unique()) + .get(); + chunk_loading_pool_.Enqueue( + [this, context](std::stop_token stop_token) { + OutputWorker(stop_token, context); + }); + } + + // Start caching workers if configured. + if (cachehit_output_queue_.has_value()) { + for (size_t i = 0; i < caching_pool_.num_threads(); ++i) { + auto* context = + caching_thread_contexts_ + .emplace_back(std::make_unique()) + .get(); + caching_pool_.Enqueue([this, context](std::stop_token stop_token) { + CachingWorker(stop_token, context); + }); + } + } + } catch (const QueueClosedException&) { + LOG(INFO) << "ShufflingChunkPool initialization interrupted, input " + "queue closed."; + output_queue()->Close(); + } catch (const std::exception& e) { + LOG(ERROR) << "ShufflingChunkPool initialization failed: " << e.what(); + output_queue()->Close(); + } + }); +} + +void ShufflingChunkPool::Stop() { + if (stop_source_.stop_requested()) return; + + LOG(INFO) << "Stopping ShufflingChunkPool."; + stop_source_.request_stop(); + if (initialization_thread_.joinable()) { + initialization_thread_.request_stop(); + initialization_thread_.join(); + } + + source_ingestion_pool_.Shutdown(); + chunk_loading_pool_.Shutdown(); + if (cachehit_output_queue_) caching_pool_.Shutdown(); + output_queue()->Close(); + if (cachehit_output_queue_) cachehit_output_queue_->Close(); + LOG(INFO) << "ShufflingChunkPool stopped."; +} + +std::vector> +ShufflingChunkPool::InitializeChunkSources() { + std::vector> uninitialized_sources; + + // Read from input queue until kInitialScanComplete. + while (true) { + auto chunk_source_with_phase = input_queue()->Get(); + + if (chunk_source_with_phase.message_type == + FilePathProvider::MessageType::kInitialScanComplete) { + LOG(INFO) + << "ShufflingChunkPool received initial scan completion marker."; + break; + } + + if (chunk_source_with_phase.message_type == + FilePathProvider::MessageType::kFile) { + // Add ChunkSource to uninitialized sources. + uninitialized_sources.push_back( + std::move(chunk_source_with_phase.source)); + } + } + + LOG(INFO) << "ShufflingChunkPool initial directory walk produced " + << uninitialized_sources.size() << " chunk source candidate(s)."; + + // Sort in descending order (newest first). + std::sort(uninitialized_sources.begin(), uninitialized_sources.end(), + [](const auto& a, const auto& b) { + return a->GetChunkSortKey() > b->GetChunkSortKey(); + }); + std::atomic total_chunks = 0; + size_t sources_to_keep = 0; + + // Process sources sequentially until we have enough chunks. + std::string current_anchor; + { + absl::MutexLock lock(&anchor_mutex_); + current_anchor = anchor_; + } + + for (auto& source : uninitialized_sources) { + if (output_queue()->IsClosed()) { + LOG(INFO) << "Output queue closed, stopping source ingestion."; + break; + } + if (total_chunks >= chunk_pool_size_) break; + + // Count chunks immediately; constructors have already prepared metadata. + const size_t chunk_count = source->GetChunkCount(); + total_chunks += chunk_count; + + // Count chunks since anchor during initial load. + if (source->GetChunkSortKey() > current_anchor) { + chunks_since_anchor_ += chunk_count; + } + + LOG_EVERY_N_SEC(INFO, 4) << "Loaded so far: " << total_chunks.load() + << "; new: " << chunks_since_anchor_; + ++sources_to_keep; + } + + LOG(INFO) << "ShufflingChunkPool indexed " << total_chunks.load() + << " chunk(s) across " << sources_to_keep + << " source(s) during startup."; + + if (total_chunks < chunk_pool_size_ && !output_queue()->IsClosed()) { + LOG(ERROR) << "ShufflingChunkPool startup chunk requirement not met: " + << total_chunks.load() << " < " << chunk_pool_size_; + } + + // Trim the vector to only keep the sources we need. + uninitialized_sources.resize(sources_to_keep); + return uninitialized_sources; +} + +void ShufflingChunkPool::ProcessInputFiles( + std::vector> uninitialized_sources) { + // Initialize chunk sources from the initial scan. + size_t initial_window_sources = 0; + size_t initial_total_chunks = 0; + { + absl::MutexLock lock(&chunk_sources_mutex_); + size_t start_chunk_index = 0; + // Newest sources first, so we add in reverse order. + std::for_each(uninitialized_sources.rbegin(), uninitialized_sources.rend(), + [this, &start_chunk_index](auto& source) { + const size_t count = source->GetChunkCount(); + auto item = std::make_shared(); + item->start_chunk_index = start_chunk_index; + item->source = std::move(source); + item->use_counts = std::vector(count, 0); + item->weight = std::vector(count, -1.0f); + item->cache = std::vector>( + cachehit_output_queue_.has_value() ? count : 0); + chunk_sources_.push_back(std::move(item)); + start_chunk_index += + chunk_sources_.back()->source->GetChunkCount(); + }); + + // Initialize stream shuffler with the initial bounds. + if (!chunk_sources_.empty()) { + size_t total_chunks = chunk_sources_.back()->start_chunk_index + + chunk_sources_.back()->source->GetChunkCount(); + // Set bounds to provide the last chunk_pool_size_ chunks. + size_t lower_bound = + total_chunks > chunk_pool_size_ ? total_chunks - chunk_pool_size_ : 0; + stream_shuffler_.SetLowerBound(lower_bound); + stream_shuffler_.SetUpperBound(total_chunks); + initial_total_chunks = total_chunks; + } + initial_window_sources = chunk_sources_.size(); + } + + LOG(INFO) << "ShufflingChunkPool initial window ready with " + << initial_window_sources << " source(s) totaling " + << initial_total_chunks << " chunk(s)."; + + // Log anchor and sources after initial scan completion. + { + absl::MutexLock anchor_lock(&anchor_mutex_); + LOG(INFO) << "Current anchor: '" << anchor_ << "'"; + + absl::MutexLock sources_lock(&chunk_sources_mutex_); + std::vector> sources_after_anchor; + for (const auto& item : chunk_sources_) { + if (item->source->GetChunkSortKey() > anchor_) { + sources_after_anchor.push_back(item); + } + } + + LOG(INFO) << sources_after_anchor.size() + << " chunk source(s) after anchor, " << chunks_since_anchor_ + << " total chunks since anchor"; + + const size_t to_log = std::min(sources_after_anchor.size(), size_t(20)); + for (size_t i = 0; i < to_log; ++i) { + LOG(INFO) << " Source [" << (i + 1) << "/" << sources_after_anchor.size() + << "]: key='" + << sources_after_anchor[i]->source->GetChunkSortKey() + << "', chunks=" + << sources_after_anchor[i]->source->GetChunkCount(); + } + } + + if (initial_total_chunks == 0) { + throw std::runtime_error( + "ShufflingChunkPool requires at least one chunk during startup."); + } +} + +void ShufflingChunkPool::SourceIngestionWorker( + std::stop_token stop_token, SourceIngestionThreadContext* context) { + try { + while (true) { + auto chunk_source_with_phase = [&]() { + LoadMetricPauser pauser(context->load_metric_updater); + return input_queue()->Get(stop_token); + }(); + + if (chunk_source_with_phase.message_type == + FilePathProvider::MessageType::kFile) { + // Ingest the new chunk source. + auto source = std::move(chunk_source_with_phase.source); + size_t chunk_count = source->GetChunkCount(); + absl::MutexLock lock(&chunk_sources_mutex_); + chunks_since_anchor_ += chunk_count; + AddNewChunkSource(std::move(source)); + } + } + } catch (const QueueClosedException&) { + LOG(INFO) << "SourceIngestionWorker stopping, queue closed."; + } catch (const QueueRequestCancelled&) { + LOG(INFO) << "SourceIngestionWorker stopping, request cancelled."; + } +} + +void ShufflingChunkPool::OutputWorker(std::stop_token stop_token, + ChunkLoadingThreadContext* context) { + // Create a local producer for this worker + auto primary_producer = output_queue()->CreateProducer(); + std::optionalCreateProducer())> + cachehit_producer; + if (cachehit_output_queue_.has_value()) { + cachehit_producer.emplace(cachehit_output_queue_->CreateProducer()); + } + + try { + while (true) { + auto result = GetNextChunkData(); + if (!result) { + if (output_queue()->IsClosed()) break; + continue; + } + LoadMetricPauser pauser(context->load_metric_updater); + if (std::holds_alternative(*result)) { + primary_producer.Put(std::move(std::get(*result)), + stop_token); + } else { + cachehit_producer->Put(std::move(std::get(*result)), + stop_token); + } + } + } catch (const QueueClosedException&) { + LOG(INFO) << "OutputWorker stopping, queue closed."; + } catch (const QueueRequestCancelled&) { + LOG(INFO) << "OutputWorker stopping, request cancelled."; + } catch (const std::exception& e) { + LOG(FATAL) << "OutputWorker encountered an error: " << e.what(); + } +} + +void ShufflingChunkPool::CachingWorker(std::stop_token stop_token, + CachingThreadContext* context) { + constexpr double kTheta = 0.99; + double reminder = 0.0; + double exponential_avg_probability = 1.0; + try { + while (true) { + auto cache_request = [&]() { + LoadMetricPauser pauser(context->load_metric_updater); + return cache_request_queue_->Get(stop_token); + }(); + + std::shared_ptr source_item; + float max_weight = 0.0f; + size_t local_index = 0; + { + absl::MutexLock lock(&chunk_sources_mutex_); + + // Find the chunk source containing this global index. + auto it = absl::c_lower_bound( + chunk_sources_, cache_request.global_index, + [](const auto& item, size_t chunk_idx) { + return item->start_chunk_index + item->source->GetChunkCount() <= + chunk_idx; + }); + + if (it == chunk_sources_.end() || + cache_request.global_index < (*it)->start_chunk_index) { + chunk_source_not_found_.fetch_add(1, std::memory_order_acq_rel); + continue; + } + + source_item = *it; + max_weight = max_weight_; + local_index = cache_request.global_index - source_item->start_chunk_index; + } + + absl::MutexLock item_lock(&source_item->mutex); + assert(local_index < source_item->use_counts.size()); + + // Check use_count match. + if (source_item->use_counts[local_index] != cache_request.next_use) { + mismatched_use_counts_.fetch_add(1, std::memory_order_acq_rel); + continue; + } + + // Compute how many positions to cache. + const float weight = source_item->weight[local_index]; + assert(weight >= 0.0f); + const double probability = ComputeHanseProbability(weight, max_weight); + exponential_avg_probability = + exponential_avg_probability * (1.0 - kTheta) + probability * kTheta; + const double n = (probability * config_.position_cache_size() / + chunk_pool_size_ / exponential_avg_probability) + + reminder; + reminder = n - std::floor(n); + const size_t positions_to_cache = static_cast(std::floor(n)); + // Traverse and extend the cache chain. + std::unique_ptr* current = &source_item->cache[local_index]; + for (size_t i = 0; i < positions_to_cache; ++i) { + if (*current) { + current = &(*current)->next; + continue; + } + if (i >= cache_request.items.size()) break; + auto node = std::make_unique(); + node->frame = cache_request.items[i]; + *current = std::move(node); + current = &(*current)->next; + newly_cached_.fetch_add(1, std::memory_order_acq_rel); + cached_positions_.fetch_add(1, std::memory_order_acq_rel); + } + + const size_t dropped = + cache_request.items.size() > positions_to_cache + ? cache_request.items.size() - positions_to_cache + : 0; + dropped_cache_positions_.fetch_add(dropped, std::memory_order_acq_rel); + } + } catch (const QueueClosedException&) { + LOG(INFO) << "CachingWorker stopping, queue closed."; + } catch (const QueueRequestCancelled&) { + LOG(INFO) << "CachingWorker stopping, request cancelled."; + } +} + +struct ShufflingChunkPool::ChunkData { + std::vector data; + std::string sort_key; + size_t local_index = 0; + size_t global_index = 0; + uint32_t use_count = 0; + std::shared_ptr source_item; +}; + +std::optional> +ShufflingChunkPool::GetNextChunkData() { + while (true) { + ChunkData chunk_data; + const ChunkStatus status = GetChunkInfo(chunk_data); + if (status == ChunkStatus::kEnd) return std::nullopt; + if (status == ChunkStatus::kRetry) continue; + + const bool hanse_enabled = config_.hanse_sampling_threshold() > 0; + if (hanse_enabled && !HanseAccept(chunk_data)) continue; + + { + absl::MutexLock lock(&chunk_data.source_item->mutex); + + assert(chunk_data.source_item->use_counts.size() > chunk_data.local_index); + chunk_data.use_count = + chunk_data.source_item->use_counts[chunk_data.local_index]++; + + if (cachehit_output_queue_.has_value()) { + auto& cache_chain = chunk_data.source_item->cache[chunk_data.local_index]; + if (cache_chain) { + cache_hits_.fetch_add(1, std::memory_order_acq_rel); + cached_positions_.fetch_sub(1, std::memory_order_acq_rel); + FrameType cached_frame = cache_chain->frame; + cache_chain = std::move(cache_chain->next); + return cached_frame; + } + cache_misses_.fetch_add(1, std::memory_order_acq_rel); + } + } + + if (chunk_data.data.empty() && !LoadChunkData(chunk_data)) continue; + + TrainingChunk chunk; + chunk.sort_key = std::move(chunk_data.sort_key); + chunk.index_within_sort_key = chunk_data.local_index; + chunk.use_count = chunk_data.use_count; + chunk.global_index = chunk_data.global_index; + chunk.frames = std::move(chunk_data.data); + + return chunk; + } +} + +bool ShufflingChunkPool::LoadChunkData(ChunkData& chunk_data) { + std::optional> data = + chunk_data.source_item->source->GetChunkData(chunk_data.local_index); + + if (!data || data->empty()) { + absl::MutexLock lock(&chunk_data.source_item->mutex); + chunk_data.source_item->dropped_chunks.insert(chunk_data.local_index); + dropped_chunks_metric_.fetch_add(1, std::memory_order_acq_rel); + return false; + } + + chunk_data.data = std::move(*data); + return true; +} + +ShufflingChunkPool::ChunkStatus ShufflingChunkPool::GetChunkInfo( + ChunkData& out_chunk_data) { + std::shared_ptr source_item; + { + absl::MutexLock lock(&chunk_sources_mutex_); + std::optional chunk_index = stream_shuffler_.GetNextItem(); + + if (!chunk_index && !chunk_sources_.empty()) { + size_t total_chunks = chunk_sources_.back()->start_chunk_index + + chunk_sources_.back()->source->GetChunkCount(); + size_t lower_bound = total_chunks > chunk_pool_size_ + ? total_chunks - chunk_pool_size_ + : chunk_sources_.front()->start_chunk_index; + stream_shuffler_.Reset(lower_bound, total_chunks); + reshuffles_.fetch_add(1, std::memory_order_acq_rel); + chunk_index = stream_shuffler_.GetNextItem(); + } + + if (!chunk_index) return ChunkStatus::kEnd; + + auto it = + absl::c_lower_bound(chunk_sources_, *chunk_index, + [](const auto& item, size_t chunk_idx) { + return item->start_chunk_index + + item->source->GetChunkCount() <= + chunk_idx; + }); + + if (ABSL_PREDICT_FALSE(it == chunk_sources_.end() || + *chunk_index < (*it)->start_chunk_index)) { + LOG(WARNING) << "Chunk index " << *chunk_index + << " out of range for available chunk sources."; + return ChunkStatus::kRetry; + } + + source_item = *it; + out_chunk_data.local_index = *chunk_index - source_item->start_chunk_index; + out_chunk_data.sort_key = source_item->source->GetChunkSortKey(); + out_chunk_data.global_index = *chunk_index; + } + + { + absl::MutexLock lock(&source_item->mutex); + if (source_item->dropped_chunks.contains(out_chunk_data.local_index)) { + return ChunkStatus::kRetry; + } + } + + out_chunk_data.source_item = std::move(source_item); + return ChunkStatus::kOk; +} + +double ShufflingChunkPool::ComputeHanseProbability(float weight, + float max_weight) const { + if (max_weight <= 0.0f) return 1.0; + return std::pow(weight / max_weight, config_.hanse_sampling_gamma()); +} + +float ShufflingChunkPool::ComputeChunkWeight( + absl::Span frames) const { + return absl::c_accumulate(frames, 0.0f, [this](float sum, const auto& frame) { + return sum + + ComputePositionSamplingWeight(frame, config_.position_sampling()); + }); +} + +bool ShufflingChunkPool::HanseAccept(ChunkData& chunk_data) { + assert(chunk_data.source_item); + float weight = -1.0f; + { + absl::MutexLock lock(&chunk_data.source_item->mutex); + assert(chunk_data.source_item->weight.size() > chunk_data.local_index); + weight = chunk_data.source_item->weight[chunk_data.local_index]; + } + + if (weight < 0.0f) { + if (chunk_data.data.empty() && !LoadChunkData(chunk_data)) return false; + weight = ComputeChunkWeight(chunk_data.data); + { + absl::MutexLock pool_lock(&chunk_sources_mutex_); + max_weight_ = std::max(max_weight_, weight); + AddSample(chunk_weight_stats_, static_cast(weight)); + } + { + absl::MutexLock lock(&chunk_data.source_item->mutex); + const float cached_weight = + chunk_data.source_item->weight[chunk_data.local_index]; + if (cached_weight < 0.0f) { + chunk_data.source_item->weight[chunk_data.local_index] = weight; + } else { + weight = cached_weight; + } + } + hanse_cache_misses_.fetch_add(1, std::memory_order_acq_rel); + } else { + hanse_cache_hits_.fetch_add(1, std::memory_order_acq_rel); + } + + float max_weight = 0.0f; + { + absl::MutexLock lock(&chunk_sources_mutex_); + max_weight = max_weight_; + } + const double p = ComputeHanseProbability(weight, max_weight); + const double u = absl::Uniform(bitgen_, 0.0, 1.0); + if (u >= p) { + hanse_rejected_.fetch_add(1, std::memory_order_acq_rel); + return false; + } + return true; +} + +void ShufflingChunkPool::AddNewChunkSource(std::unique_ptr source) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(chunk_sources_mutex_) { + // Add new chunk source to the end of the deque. + size_t old_upper_bound = 0; + if (!chunk_sources_.empty()) { + const auto& last_source = chunk_sources_.back(); + old_upper_bound = + last_source->start_chunk_index + last_source->source->GetChunkCount(); + } + + size_t count = source->GetChunkCount(); + auto item = std::make_shared(); + item->start_chunk_index = old_upper_bound; + item->source = std::move(source); + item->use_counts = std::vector(count, 0); + item->weight = std::vector(count, -1.0f); + item->cache = + std::vector>(cachehit_output_queue_.has_value() + ? count + : 0); + chunk_sources_.push_back(std::move(item)); + + // Calculate current window bounds. + size_t new_upper_bound = chunk_sources_.back()->start_chunk_index + + chunk_sources_.back()->source->GetChunkCount(); + + // Remove old chunks if window exceeds chunk_pool_size_. + while (!chunk_sources_.empty() && chunk_sources_.size() > 1) { + size_t window_start = chunk_sources_.front()->start_chunk_index + + chunk_sources_.front()->source->GetChunkCount(); + size_t window_size = new_upper_bound - window_start; + + if (window_size < chunk_pool_size_) break; + + // Count cached positions in the evicted source. + if (cachehit_output_queue_.has_value()) { + size_t evicted_cached = 0; + absl::MutexLock item_lock(&chunk_sources_.front()->mutex); + for (const auto& cache_chain : chunk_sources_.front()->cache) { + const CacheNode* node = cache_chain.get(); + while (node) { + ++evicted_cached; + node = node->next.get(); + } + } + cached_positions_.fetch_sub(evicted_cached, std::memory_order_acq_rel); + } + + // Remove the oldest chunk source (front of deque). + chunk_sources_.pop_front(); + } + + // Update stream shuffler bounds with the sliding window. + size_t window_start = chunk_sources_.front()->start_chunk_index; + size_t new_lower_bound = new_upper_bound > chunk_pool_size_ + ? new_upper_bound - chunk_pool_size_ + : window_start; + stream_shuffler_.SetUpperBound(new_upper_bound); + stream_shuffler_.SetLowerBound(new_lower_bound); +} + +StageMetricProto ShufflingChunkPool::FlushMetrics() { + StageMetricProto stage_metric; + // Aggregate source ingestion load metrics from all ingestion threads. + LoadMetricProto ingestion_load; + ingestion_load.set_name("source_ingestion"); + for (const auto& context : source_ingestion_thread_contexts_) { + UpdateFrom(ingestion_load, context->load_metric_updater.FlushMetrics()); + } + *stage_metric.add_load_metrics() = std::move(ingestion_load); + + // Aggregate chunk loading load metrics from all chunk loading threads. + LoadMetricProto chunk_loading_load; + chunk_loading_load.set_name("chunk_loading"); + for (const auto& context : chunk_loading_thread_contexts_) { + UpdateFrom(chunk_loading_load, context->load_metric_updater.FlushMetrics()); + } + *stage_metric.add_load_metrics() = std::move(chunk_loading_load); + + // Get chunk sources statistics and pool state. + { + absl::MutexLock lock(&chunk_sources_mutex_); + auto* chunk_sources_metric = stage_metric.add_gauge_metrics(); + chunk_sources_metric->set_name("chunk_sources"); + chunk_sources_metric->set_value( + static_cast(chunk_sources_.size())); + + size_t upper = 0; + size_t current = 0; + if (!chunk_sources_.empty()) { + const auto& first = chunk_sources_.front(); + const auto& last = chunk_sources_.back(); + upper = last->start_chunk_index + last->source->GetChunkCount(); + current = upper - first->start_chunk_index; + } + + auto* current_chunks_metric = stage_metric.add_gauge_metrics(); + current_chunks_metric->set_name("chunks_current"); + current_chunks_metric->set_value(static_cast(current)); + current_chunks_metric->set_capacity( + static_cast(chunk_pool_size_)); + + auto* total_chunks_metric = stage_metric.add_gauge_metrics(); + total_chunks_metric->set_name("chunks_total"); + total_chunks_metric->set_value(static_cast(upper)); + } + + // Get anchor-related metrics. + { + absl::MutexLock lock(&anchor_mutex_); + auto* chunks_since_anchor_metric = stage_metric.add_gauge_metrics(); + chunks_since_anchor_metric->set_name("chunks_since_anchor"); + chunks_since_anchor_metric->set_value(chunks_since_anchor_); + stage_metric.set_anchor(anchor_); + } + + auto* dropped_metric = stage_metric.add_count_metrics(); + dropped_metric->set_name("dropped"); + dropped_metric->set_count( + dropped_chunks_metric_.exchange(0, std::memory_order_acq_rel)); + + // Hanse sampling and shuffler metrics. + { + auto* hits = stage_metric.add_count_metrics(); + hits->set_name("hanse_cache_hits"); + hits->set_count(hanse_cache_hits_.exchange(0, std::memory_order_acq_rel)); + + auto* misses = stage_metric.add_count_metrics(); + misses->set_name("hanse_cache_misses"); + misses->set_count( + hanse_cache_misses_.exchange(0, std::memory_order_acq_rel)); + + auto* rejected = stage_metric.add_count_metrics(); + rejected->set_name("hanse_rejected"); + rejected->set_count(hanse_rejected_.exchange(0, std::memory_order_acq_rel)); + + auto* resh = stage_metric.add_count_metrics(); + resh->set_name("reshuffles"); + resh->set_count(reshuffles_.exchange(0, std::memory_order_acq_rel)); + } + + // Position cache metrics. + if (cachehit_output_queue_.has_value()) { + LoadMetricProto caching_load; + caching_load.set_name("caching"); + for (const auto& context : caching_thread_contexts_) { + UpdateFrom(caching_load, context->load_metric_updater.FlushMetrics()); + } + *stage_metric.add_load_metrics() = std::move(caching_load); + + auto* cache_hits = stage_metric.add_count_metrics(); + cache_hits->set_name("cache_hits"); + cache_hits->set_count(cache_hits_.exchange(0, std::memory_order_acq_rel)); + + auto* cache_misses = stage_metric.add_count_metrics(); + cache_misses->set_name("cache_misses"); + cache_misses->set_count( + cache_misses_.exchange(0, std::memory_order_acq_rel)); + + auto* mismatched = stage_metric.add_count_metrics(); + mismatched->set_name("mismatched_use_counts"); + mismatched->set_count( + mismatched_use_counts_.exchange(0, std::memory_order_acq_rel)); + + auto* newly_cached = stage_metric.add_count_metrics(); + newly_cached->set_name("newly_cached"); + newly_cached->set_count( + newly_cached_.exchange(0, std::memory_order_acq_rel)); + + auto* dropped = stage_metric.add_count_metrics(); + dropped->set_name("dropped_cache_positions"); + dropped->set_count( + dropped_cache_positions_.exchange(0, std::memory_order_acq_rel)); + + auto* not_found = stage_metric.add_count_metrics(); + not_found->set_name("chunk_source_not_found"); + not_found->set_count( + chunk_source_not_found_.exchange(0, std::memory_order_acq_rel)); + + auto* cached = stage_metric.add_gauge_metrics(); + cached->set_name("cached_positions"); + cached->set_value(cached_positions_.load(std::memory_order_acquire)); + cached->set_capacity(config_.position_cache_size()); + } + + { + absl::MutexLock lock(&chunk_sources_mutex_); + if (chunk_weight_stats_.count() > 0) { + chunk_weight_stats_.set_name("chunk_weight"); + UpdateFrom(*stage_metric.add_statistics_metrics(), chunk_weight_stats_); + } + chunk_weight_stats_.Clear(); + } + + *stage_metric.add_queue_metrics() = + MetricsFromQueue(primary_output_name_, *output_queue()); + if (cachehit_output_queue_.has_value()) { + *stage_metric.add_queue_metrics() = + MetricsFromQueue(*cachehit_output_name_, *cachehit_output_queue_); + } + return stage_metric; +} + +std::pair ShufflingChunkPool::ResetAnchor() { + absl::MutexLock anchor_lock(&anchor_mutex_); + absl::MutexLock sources_lock(&chunk_sources_mutex_); + + if (chunk_sources_.empty()) { + int previous_count = chunks_since_anchor_.exchange(0); + return {anchor_, previous_count}; + } + + anchor_ = chunk_sources_.back()->source->GetChunkSortKey(); + int previous_count = chunks_since_anchor_.exchange(0); + return {anchor_, previous_count}; +} + +int ShufflingChunkPool::ChunksSinceAnchor() { return chunks_since_anchor_; } + +std::string ShufflingChunkPool::CurrentAnchor() { + absl::MutexLock lock(&anchor_mutex_); + return anchor_; +} + +void ShufflingChunkPool::SetAnchor(std::string_view anchor) { + absl::MutexLock lock(&anchor_mutex_); + anchor_ = anchor; +} + +std::optional ShufflingChunkPool::Control( + const StageControlRequest& request) { + if (!request.has_chunk_pool_request()) { + return std::nullopt; + } + + const auto& chunk_request = request.chunk_pool_request(); + StageControlResponse response; + auto* chunk_response = response.mutable_chunk_pool_response(); + + if (chunk_request.reset_chunk_anchor()) { + auto [anchor, chunks] = ResetAnchor(); + chunk_response->set_chunk_anchor(anchor); + chunk_response->set_chunks_since_anchor(chunks); + return response; + } + + if (chunk_request.has_set_chunk_anchor()) { + SetAnchor(chunk_request.set_chunk_anchor()); + chunk_response->set_chunk_anchor(chunk_request.set_chunk_anchor()); + chunk_response->set_chunks_since_anchor(ChunksSinceAnchor()); + return response; + } + + chunk_response->set_chunk_anchor(CurrentAnchor()); + chunk_response->set_chunks_since_anchor(ChunksSinceAnchor()); + return response; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/shuffling_chunk_pool.h b/csrc/loader/stages/shuffling_chunk_pool.h new file mode 100644 index 00000000..d499aac8 --- /dev/null +++ b/csrc/loader/stages/shuffling_chunk_pool.h @@ -0,0 +1,163 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/flat_hash_set.h" +#include "absl/random/random.h" +#include "absl/synchronization/mutex.h" +#include "loader/chunk_source/chunk_source.h" +#include "loader/data_loader_metrics.h" +#include "loader/stages/chunk_source_loader.h" +#include "loader/stages/stage.h" +#include "loader/stages/training_chunk.h" +#include "proto/data_loader_config.pb.h" +#include "proto/stage_control.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/metrics/load_metric.h" +#include "utils/queue.h" +#include "utils/stream_shuffler.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +class ShufflingChunkPool : public Stage { + public: + explicit ShufflingChunkPool(const ShufflingChunkPoolConfig& config); + ~ShufflingChunkPool(); + + void Start() override; + void Stop() override; + void SetInputs(absl::Span inputs) override; + QueueBase* GetOutput(std::string_view name) override; + + StageMetricProto FlushMetrics() override; + + std::optional Control( + const StageControlRequest& request) override; + + // Anchor management methods for tracking chunks since a specific point. + std::pair ResetAnchor(); + int ChunksSinceAnchor(); + std::string CurrentAnchor(); + void SetAnchor(std::string_view anchor); + + Queue* input_queue() { return primary_input_queue_; } + Queue* output_queue() { return &primary_output_queue_; } + + private: + struct CacheNode { + FrameType frame; + std::unique_ptr next; + }; + + struct ChunkSourceItem { + mutable absl::Mutex mutex; + size_t start_chunk_index; + std::unique_ptr source; + absl::flat_hash_set dropped_chunks ABSL_GUARDED_BY(mutex); + // Per-chunk counters and cached weights. + std::vector use_counts ABSL_GUARDED_BY(mutex); + std::vector weight ABSL_GUARDED_BY(mutex); + std::vector> cache ABSL_GUARDED_BY(mutex); + }; + + struct SourceIngestionThreadContext { + LoadMetricUpdater load_metric_updater; + }; + + struct ChunkLoadingThreadContext { + LoadMetricUpdater load_metric_updater; + }; + + struct CachingThreadContext { + LoadMetricUpdater load_metric_updater; + }; + + std::vector> InitializeChunkSources(); + void ProcessInputFiles( + std::vector> uninitialized_sources); + void SourceIngestionWorker(std::stop_token stop_token, + SourceIngestionThreadContext* context); + void OutputWorker(std::stop_token stop_token, + ChunkLoadingThreadContext* context); + void CachingWorker(std::stop_token stop_token, CachingThreadContext* context); + void AddNewChunkSource(std::unique_ptr source) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(chunk_sources_mutex_); + std::optional> GetNextChunkData(); + + enum class ChunkStatus { kOk, kRetry, kEnd }; + struct ChunkData; + + ChunkStatus GetChunkInfo(ChunkData& out_chunk_data); + bool LoadChunkData(ChunkData& chunk_data); + bool HanseAccept(ChunkData& chunk_data); + float ComputeChunkWeight(absl::Span frames) const; + double ComputeHanseProbability(float weight, float max_weight) const; + + Queue* primary_input_queue_ = nullptr; + Queue* cache_request_queue_ = nullptr; + std::string primary_output_name_; + Queue primary_output_queue_; + std::optional cachehit_output_name_; + std::optional> cachehit_output_queue_; + + const size_t chunk_pool_size_; + const ShufflingChunkPoolConfig config_; + // stop_source_ must be declared before ThreadPools that reference it. + std::stop_source stop_source_; + ThreadPool source_ingestion_pool_; + ThreadPool chunk_loading_pool_; + ThreadPool caching_pool_; + + std::atomic dropped_chunks_metric_{0}; + + absl::Mutex chunk_sources_mutex_; + std::deque> chunk_sources_ + ABSL_GUARDED_BY(chunk_sources_mutex_); + StreamShuffler stream_shuffler_ ABSL_GUARDED_BY(chunk_sources_mutex_); + float max_weight_ ABSL_GUARDED_BY(chunk_sources_mutex_) = 0.0f; + std::jthread initialization_thread_; + std::vector> + source_ingestion_thread_contexts_; + std::vector> + chunk_loading_thread_contexts_; + std::vector> caching_thread_contexts_; + + // Anchor-related members for tracking chunks since a specific point. + absl::Mutex anchor_mutex_; + std::string anchor_ ABSL_GUARDED_BY(anchor_mutex_); + std::atomic chunks_since_anchor_{0}; + + // Thread-local RNG for Hanse sampling. + static thread_local absl::BitGen bitgen_; + + // Metrics counters. + std::atomic hanse_cache_hits_{0}; + std::atomic hanse_cache_misses_{0}; + std::atomic hanse_rejected_{0}; + std::atomic reshuffles_{0}; + std::atomic cache_hits_{0}; + std::atomic cache_misses_{0}; + std::atomic mismatched_use_counts_{0}; + std::atomic newly_cached_{0}; + std::atomic dropped_cache_positions_{0}; + std::atomic chunk_source_not_found_{0}; + std::atomic cached_positions_{0}; + + StatisticsProtoDouble chunk_weight_stats_ + ABSL_GUARDED_BY(chunk_sources_mutex_); +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/shuffling_chunk_pool_test.cc b/csrc/loader/stages/shuffling_chunk_pool_test.cc new file mode 100644 index 00000000..32e58557 --- /dev/null +++ b/csrc/loader/stages/shuffling_chunk_pool_test.cc @@ -0,0 +1,844 @@ +// ABOUTME: Comprehensive unit tests for the ShufflingChunkPool class +// ABOUTME: Tests chunk source management, output workers, and dynamic windowing + +#include "loader/stages/shuffling_chunk_pool.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "loader/stages/training_chunk.h" + +namespace lczero { +namespace training { + +namespace { + +template +class PassthroughStage : public Stage { + public: + explicit PassthroughStage(Queue* queue) : queue_(queue) {} + + void Start() override {} + void Stop() override {} + StageMetricProto FlushMetrics() override { return StageMetricProto(); } + QueueBase* GetOutput(std::string_view name = "") override { + (void)name; + return queue_; + } + void SetInputs(absl::Span inputs) override { + if (!inputs.empty()) { + throw std::runtime_error("PassthroughStage expects no inputs"); + } + } + + private: + Queue* queue_; +}; + +} // namespace + +// Mock ChunkSource for testing +class MockChunkSource : public ChunkSource { + public: + MockChunkSource(const std::string& sort_key, size_t chunk_count) + : sort_key_(sort_key), chunk_count_(chunk_count) {} + + std::string GetChunkSortKey() const override { return sort_key_; } + size_t GetChunkCount() const override { return chunk_count_; } + + std::optional> GetChunkData(size_t index) override { + if (index >= chunk_count_) { + throw std::out_of_range("Chunk index out of range"); + } + FrameType frame{}; + frame.version = static_cast(index); + frame.input_format = 3; + return std::vector{frame}; + } + + private: + std::string sort_key_; + size_t chunk_count_; +}; + +class InvalidChunkSource : public ChunkSource { + public: + explicit InvalidChunkSource(std::string sort_key) + : sort_key_(std::move(sort_key)) {} + + std::string GetChunkSortKey() const override { return sort_key_; } + size_t GetChunkCount() const override { return 2; } + + std::optional> GetChunkData(size_t index) override { + if (index >= 2) { + throw std::out_of_range("Chunk index out of range"); + } + if (index == 0) { + return std::nullopt; + } + FrameType frame{}; + frame.version = 42; + return std::vector{frame}; + } + + private: + std::string sort_key_; +}; + +class ShufflingChunkPoolTest : public ::testing::Test { + protected: + void SetUp() override { + input_queue_ = std::make_unique>(100); + input_producer_ = std::make_unique::Producer>( + input_queue_->CreateProducer()); + } + + void TearDown() override { + // Close the producer to close the queue + if (input_producer_) input_producer_.reset(); + } + + // Helper to add a mock chunk source to the input queue + void AddMockChunkSourceToQueue(const std::string& sort_key, + size_t chunk_count, + FilePathProvider::MessageType message_type = + FilePathProvider::MessageType::kFile) { + ChunkSourceWithPhase item; + item.source = std::make_unique(sort_key, chunk_count); + item.message_type = message_type; + input_producer_->Put(std::move(item)); + } + + void MarkInitialScanComplete() { + ChunkSourceWithPhase item; + item.source = nullptr; // No source for completion marker + item.message_type = FilePathProvider::MessageType::kInitialScanComplete; + input_producer_->Put(std::move(item)); + } + + void CloseInputQueue() { + if (input_producer_) input_producer_.reset(); + } + + ShufflingChunkPoolConfig MakeConfig(int chunk_pool_size, + int source_ingestion_threads = 1, + int loading_threads = 1, + int queue_capacity = 100) const { + ShufflingChunkPoolConfig config; + config.set_chunk_pool_size(chunk_pool_size); + config.set_source_ingestion_threads(source_ingestion_threads); + config.set_chunk_loading_threads(loading_threads); + config.mutable_output()->set_queue_capacity(queue_capacity); + return config; + } + + std::unique_ptr> input_queue_; + std::unique_ptr::Producer> input_producer_; +}; + +TEST_F(ShufflingChunkPoolTest, ConstructorCreatesOutputQueue) { + // Add some mock chunk sources with enough chunks + AddMockChunkSourceToQueue("source1", 50); + AddMockChunkSourceToQueue("source2", 60); + MarkInitialScanComplete(); + + auto config = MakeConfig(20); + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + + auto* output_queue = shuffling_chunk_pool.output_queue(); + + // Close input queue to stop input worker from waiting + CloseInputQueue(); + + EXPECT_NE(output_queue, nullptr); + EXPECT_EQ(output_queue->Capacity(), 100); + + // Drain output queue to prevent workers from blocking + try { + while (output_queue->Size() > 0) { + output_queue->Get(); + } + } catch (const QueueClosedException&) { + // Queue closed, that's fine + } +} + +TEST_F(ShufflingChunkPoolTest, HandlesEmptyInputQueue) { + // Only mark scan complete, no chunk sources + MarkInitialScanComplete(); + + auto config = MakeConfig(20); + + // Constructor should now succeed (initialization is asynchronous) + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + // The initialization thread should handle the error case + auto* output_queue = shuffling_chunk_pool.output_queue(); + + // Give the initialization thread time to complete and discover the error + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + // Close input queue to clean up + CloseInputQueue(); + + // Output queue should exist but be closed to signal startup failure when no + // chunks were found. + EXPECT_NE(output_queue, nullptr); + EXPECT_TRUE(output_queue->IsClosed()); + EXPECT_EQ(output_queue->Size(), 0u); +} + +TEST_F(ShufflingChunkPoolTest, FlushMetricsHandlesEmptyChunkSources) { + const int chunk_pool_size = 32; + auto config = MakeConfig(chunk_pool_size); + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + + auto metrics = shuffling_chunk_pool.FlushMetrics(); + bool found_current = false; + bool found_total = false; + for (const auto& metric : metrics.gauge_metrics()) { + if (metric.name() == "chunks_current") { + found_current = true; + EXPECT_EQ(metric.value(), 0u); + EXPECT_EQ(metric.capacity(), static_cast(chunk_pool_size)); + } else if (metric.name() == "chunks_total") { + found_total = true; + EXPECT_EQ(metric.value(), 0u); + } + } + + EXPECT_TRUE(found_current) + << "FlushMetrics should emit chunks_current metric when empty."; + EXPECT_TRUE(found_total) + << "FlushMetrics should emit chunks_total metric when empty."; +} + +TEST_F(ShufflingChunkPoolTest, FlushMetricsReportsWindowAndTotalCounts) { + AddMockChunkSourceToQueue("initial", 30); + MarkInitialScanComplete(); + + const int chunk_pool_size = 20; + ShufflingChunkPool shuffling_chunk_pool(MakeConfig(chunk_pool_size)); + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + auto* output_queue = shuffling_chunk_pool.output_queue(); + output_queue->WaitForSizeAtLeast(1); + + uint64_t current_count = 0; + uint64_t total_count = 0; + uint64_t current_capacity = 0; + bool found_metrics = false; + for (int attempt = 0; attempt < 50 && !found_metrics; ++attempt) { + auto metrics = shuffling_chunk_pool.FlushMetrics(); + bool has_current = false; + bool has_total = false; + for (const auto& metric : metrics.gauge_metrics()) { + if (metric.name() == "chunks_current") { + has_current = true; + current_count = metric.value(); + current_capacity = metric.capacity(); + } else if (metric.name() == "chunks_total") { + has_total = true; + total_count = metric.value(); + } + } + if (has_current && has_total) { + found_metrics = true; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + + ASSERT_TRUE(found_metrics) + << "FlushMetrics should report both chunks_current and chunks_total."; + EXPECT_EQ(current_count, 30u); + EXPECT_EQ(current_capacity, static_cast(chunk_pool_size)); + EXPECT_EQ(total_count, 30u); + + CloseInputQueue(); +} + +TEST_F(ShufflingChunkPoolTest, ProcessesInitialScanChunkSources) { + // Create mock chunk sources with enough chunks + AddMockChunkSourceToQueue("source1", 30); + AddMockChunkSourceToQueue("source2", 40); + AddMockChunkSourceToQueue("source3", 50); + MarkInitialScanComplete(); + + auto config = MakeConfig(20); + + // Test that constructor completes and processes mock chunk sources + EXPECT_NO_THROW({ + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + // Close input queue to stop input worker from waiting + CloseInputQueue(); + + auto* output_queue = shuffling_chunk_pool.output_queue(); + EXPECT_NE(output_queue, nullptr); + }); +} + +TEST_F(ShufflingChunkPoolTest, OutputWorkerProducesChunks) { + // Create mock chunk sources + AddMockChunkSourceToQueue("source1", 10, + FilePathProvider::MessageType::kFile); + AddMockChunkSourceToQueue("source2", 15, + FilePathProvider::MessageType::kFile); + MarkInitialScanComplete(); + + auto config = MakeConfig(20); + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + // Close input queue to stop input worker from waiting + CloseInputQueue(); + + auto* output_queue = shuffling_chunk_pool.output_queue(); + + // Wait for output workers to produce at least one chunk + output_queue->WaitForSizeAtLeast(1); + + // Should have some chunks available + EXPECT_GT(output_queue->Size(), 0); + + // Get a chunk and verify it's from our mock sources + auto chunk = output_queue->Get(); + EXPECT_FALSE(chunk.frames.empty()); + EXPECT_TRUE(chunk.sort_key == "source1" || chunk.sort_key == "source2"); + EXPECT_EQ(chunk.frames.size(), 1); + EXPECT_EQ(chunk.frames.front().version, + static_cast(chunk.index_within_sort_key)); + EXPECT_EQ(chunk.use_count, 0u); +} + +TEST_F(ShufflingChunkPoolTest, DropsInvalidChunks) { + ChunkSourceWithPhase invalid_source; + invalid_source.source = + std::make_unique("invalid_source"); + invalid_source.message_type = FilePathProvider::MessageType::kFile; + input_producer_->Put(std::move(invalid_source)); + MarkInitialScanComplete(); + + auto config = MakeConfig(2, /*source_ingestion_threads=*/1, + /*loading_threads=*/1, /*queue_capacity=*/10); + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + // Close input queue to stop input worker from waiting + CloseInputQueue(); + + auto* output_queue = shuffling_chunk_pool.output_queue(); + output_queue->WaitForSizeAtLeast(1); + auto chunk = output_queue->Get(); + + EXPECT_EQ(chunk.sort_key, "invalid_source"); + EXPECT_EQ(chunk.index_within_sort_key, 1); + EXPECT_EQ(chunk.use_count, 0u); + ASSERT_EQ(chunk.frames.size(), 1); + EXPECT_EQ(chunk.frames.front().version, 42); + + uint64_t dropped_latest = 0; + bool found_dropped = false; + for (int attempt = 0; attempt < 50 && !found_dropped; ++attempt) { + auto metrics = shuffling_chunk_pool.FlushMetrics(); + for (const auto& metric : metrics.count_metrics()) { + if (metric.name() == "dropped" && metric.count() > 0) { + dropped_latest = metric.count(); + found_dropped = true; + break; + } + } + if (!found_dropped) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + } + ASSERT_TRUE(found_dropped) << "dropped chunk metrics should be reported"; + EXPECT_GE(dropped_latest, 1u); +} + +TEST_F(ShufflingChunkPoolTest, NewChunkSourceProcessing) { + // Start with initial scan and one chunk source - use enough chunks to satisfy + // window + AddMockChunkSourceToQueue("initial", 120); // More chunks than window + MarkInitialScanComplete(); + + auto config = MakeConfig(20); + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + // Verify chunks are being produced from initial sources + auto* output_queue = shuffling_chunk_pool.output_queue(); + output_queue->WaitForSizeAtLeast(1); + EXPECT_NE(output_queue, nullptr); + EXPECT_GT(output_queue->Size(), 0); + + // Add a new chunk source after initialization + AddMockChunkSourceToQueue("new_source", 30, + FilePathProvider::MessageType::kFile); + + // Close input queue to stop input worker from waiting for more + CloseInputQueue(); + + // The chunk set should still be functional and continue producing chunks + // from both the initial and new sources + EXPECT_GT(output_queue->Size(), 0); +} + +TEST_F(ShufflingChunkPoolTest, ChunkWindowManagement) { + // Create more chunks than the window size + AddMockChunkSourceToQueue("source1", 30); + AddMockChunkSourceToQueue("source2", 30); + AddMockChunkSourceToQueue("source3", 30); + MarkInitialScanComplete(); + + auto config = MakeConfig(50); + + // Should only keep sources that fit in the window + EXPECT_NO_THROW({ + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + // Close input queue to stop input worker from waiting + CloseInputQueue(); + + auto* output_queue = shuffling_chunk_pool.output_queue(); + EXPECT_NE(output_queue, nullptr); + }); +} + +// Test the ShufflingChunkPoolConfig structure +TEST_F(ShufflingChunkPoolTest, ChunkSorting) { + // Add chunk sources in non-sorted order (by sort key) + AddMockChunkSourceToQueue("source_b", 20); + AddMockChunkSourceToQueue("source_a", 25); + AddMockChunkSourceToQueue("source_c", 30); + MarkInitialScanComplete(); + + auto config = MakeConfig(70); + + // ShufflingChunkPool should handle sorting internally (newest first) + EXPECT_NO_THROW({ + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + // Close input queue to stop input worker from waiting + CloseInputQueue(); + + auto* output_queue = shuffling_chunk_pool.output_queue(); + EXPECT_NE(output_queue, nullptr); + }); +} + +TEST_F(ShufflingChunkPoolTest, StreamShufflerResetWhenExhausted) { + // Create a small chunk source to quickly exhaust the shuffler + AddMockChunkSourceToQueue("source1", 3); // Only 3 chunks for faster testing + MarkInitialScanComplete(); + + auto config = MakeConfig(3, /*source_ingestion_threads=*/1, + /*loading_threads=*/1, + /*queue_capacity=*/100); // Large enough + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + auto* output_queue = shuffling_chunk_pool.output_queue(); + + // Collect chunks continuously and count total chunks received + struct ChunkRecord { + std::string sort_key; + size_t index; + uint32_t use_count; + }; + std::vector all_chunks_received; + + // Wait for and collect chunks to test shuffler reset + for (size_t i = 0; i < 8; ++i) { + output_queue->WaitForSizeAtLeast(1); + auto chunk = output_queue->Get(); + all_chunks_received.push_back( + {chunk.sort_key, chunk.index_within_sort_key, chunk.use_count}); + } + + std::set> unique_chunks; + bool seen_reuse = false; + for (const auto& record : all_chunks_received) { + unique_chunks.emplace(record.sort_key, record.index); + if (record.use_count > 0) { + seen_reuse = true; + } + } + + // Close input queue to clean up + try { + CloseInputQueue(); + } catch (const QueueClosedException&) { + // Already closed, that's fine + } + + // We should see all 3 unique chunks from our source + EXPECT_EQ(unique_chunks.size(), 3) << "Should see all unique chunks"; + + // If reset works properly, we should receive more than 3 total chunks + // (since chunks will repeat after shuffler reset) + EXPECT_GT(all_chunks_received.size(), 3) + << "Should get more than 3 chunks total due to shuffler reset, got " + << all_chunks_received.size() << " chunks"; + EXPECT_TRUE(seen_reuse) + << "Expect at least one chunk to report a reuse count"; +} + +TEST_F(ShufflingChunkPoolTest, HanseMetrics_NoRejection_CacheAndReshuffles) { + // Single chunk so we will continually reuse the same chunk. + AddMockChunkSourceToQueue("source1", 1); + MarkInitialScanComplete(); + + auto config = MakeConfig(1, /*source_ingestion_threads=*/1, + /*loading_threads=*/1, /*queue_capacity=*/100); + // Enable Hanse sampling with p == 1 to avoid rejections. + config.set_hanse_sampling_threshold(1); + + ShufflingChunkPool pool(config); + + pool.SetInputs({input_queue_.get()}); + pool.Start(); + + auto* output_queue = pool.output_queue(); + // Wait for multiple outputs to exercise cache hits and reshuffles. + output_queue->WaitForSizeAtLeast(3); + // Drain a few items. + for (int i = 0; i < 3; ++i) { + auto chunk = output_queue->Get(); + EXPECT_EQ(chunk.frames.size(), 1u); + } + + // Close input to avoid lingering. + CloseInputQueue(); + + // Flush metrics and validate Hanse counters and reshuffles. + auto metrics = pool.FlushMetrics(); + uint64_t cache_hits = 0, cache_misses = 0, rejected = 0, reshuffles = 0; + for (const auto& m : metrics.count_metrics()) { + if (m.name() == "hanse_cache_hits") cache_hits = m.count(); + if (m.name() == "hanse_cache_misses") cache_misses = m.count(); + if (m.name() == "hanse_rejected") rejected = m.count(); + if (m.name() == "reshuffles") reshuffles = m.count(); + } + + // First access computes and caches num_records => 1 miss, then hits. + EXPECT_EQ(cache_misses, 1u); + EXPECT_GE(cache_hits, 1u); + // With threshold=1 and one frame, p = 1 => no rejections. + EXPECT_EQ(rejected, 0u); + // Single chunk repeatedly consumed forces reshuffles. + EXPECT_GT(reshuffles, 0u); +} + +TEST_F(ShufflingChunkPoolTest, ExplicitClose) { + // Create chunk sources + AddMockChunkSourceToQueue("source1", 20); + AddMockChunkSourceToQueue("source2", 30); + MarkInitialScanComplete(); + + auto config = MakeConfig(40); + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + auto* output_queue = shuffling_chunk_pool.output_queue(); + + // Wait for workers to produce some chunks + output_queue->WaitForSizeAtLeast(1); + + // Verify output queue is working before close + EXPECT_GT(output_queue->Size(), 0); + + // Explicitly stop the chunk set + shuffling_chunk_pool.Stop(); + + // Drain all remaining items from the queue + while (output_queue->Size() > 0) { + output_queue->Get(); + } + + // Now the queue should be closed and empty, so Get() should throw + EXPECT_THROW(output_queue->Get(), QueueClosedException); + + CloseInputQueue(); +} + +TEST_F(ShufflingChunkPoolTest, CloseStopsOutputWorkers) { + // Create chunk sources + AddMockChunkSourceToQueue("source1", 15); + MarkInitialScanComplete(); + + auto config = MakeConfig(15, /*source_ingestion_threads=*/1, + /*loading_threads=*/2, /*queue_capacity=*/50); + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + auto* output_queue = shuffling_chunk_pool.output_queue(); + + // Wait for workers to produce chunks + output_queue->WaitForSizeAtLeast(1); + size_t chunks_before_close = output_queue->Size(); + + // Stop the chunk set + shuffling_chunk_pool.Stop(); + + // Drain any remaining chunks from the queue + try { + while (output_queue->Size() > 0) { + output_queue->Get(); + } + } catch (const QueueClosedException&) { + // Expected when queue is empty and closed + } + + // Should have had chunks before close + EXPECT_GT(chunks_before_close, 0) << "Should have had chunks before close"; + + CloseInputQueue(); +} + +TEST_F(ShufflingChunkPoolTest, CloseIsIdempotent) { + // Create chunk sources + AddMockChunkSourceToQueue("source1", 20); + MarkInitialScanComplete(); + + auto config = MakeConfig(20); + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + + // Stop multiple times - should not crash or cause issues + EXPECT_NO_THROW(shuffling_chunk_pool.Stop()); + EXPECT_NO_THROW(shuffling_chunk_pool.Stop()); + EXPECT_NO_THROW(shuffling_chunk_pool.Stop()); + + CloseInputQueue(); +} + +TEST_F(ShufflingChunkPoolTest, DestructorCallsClose) { + // Create chunk sources + AddMockChunkSourceToQueue("source1", 20); + MarkInitialScanComplete(); + + auto config = MakeConfig(20); + + // Test that destructor calls Close() and properly shuts down + { + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + auto* output_queue = shuffling_chunk_pool.output_queue(); + + // Wait for workers to produce some chunks + output_queue->WaitForSizeAtLeast(1); + EXPECT_GT(output_queue->Size(), 0); + + // Close input queue before destructor to allow threads to finish + CloseInputQueue(); + + // ShufflingChunkPool destructor should be called here, which calls Stop() + // and waits for all threads to finish + } + + // Test passes if destructor completes without hanging + // (we can't test the queue state after destruction since it's destroyed) +} + +TEST_F(ShufflingChunkPoolTest, InputQueueClosureDoesNotCloseOutputQueue) { + // Create chunk sources + AddMockChunkSourceToQueue("source1", 30); + MarkInitialScanComplete(); + + auto config = MakeConfig(30); + + ShufflingChunkPool shuffling_chunk_pool(config); + + shuffling_chunk_pool.SetInputs({input_queue_.get()}); + shuffling_chunk_pool.Start(); + auto* output_queue = shuffling_chunk_pool.output_queue(); + + // Wait for workers to produce some chunks + output_queue->WaitForSizeAtLeast(1); + EXPECT_GT(output_queue->Size(), 0); + + // Close input queue (simulating end of file discovery) + CloseInputQueue(); + + // Output queue should still be functional - workers should continue + // producing chunks from existing chunk sources + + // Should still be able to get chunks (queue not closed) + EXPECT_NO_THROW(output_queue->Get()); + + // Explicitly stop to clean up + shuffling_chunk_pool.Stop(); +} + +TEST_F(ShufflingChunkPoolTest, BasicAnchorFunctionality) { + AddMockChunkSourceToQueue("source1", 20); + MarkInitialScanComplete(); + + auto config = MakeConfig(20); + + ShufflingChunkPool pool(config); + + pool.SetInputs({input_queue_.get()}); + pool.Start(); + + // Test initial state + EXPECT_EQ(pool.ChunksSinceAnchor(), 0); + EXPECT_EQ(pool.CurrentAnchor(), ""); + + // Test SetAnchor and CurrentAnchor + pool.SetAnchor("test_anchor_key"); + EXPECT_EQ(pool.CurrentAnchor(), "test_anchor_key"); + EXPECT_EQ(pool.ChunksSinceAnchor(), 0); // Should still be 0 + + // Test setting different anchor + pool.SetAnchor("another_key"); + EXPECT_EQ(pool.CurrentAnchor(), "another_key"); + + CloseInputQueue(); +} + +TEST_F(ShufflingChunkPoolTest, ResetAnchor) { + AddMockChunkSourceToQueue("source1", 20); + MarkInitialScanComplete(); + + auto config = MakeConfig(20); + + ShufflingChunkPool pool(config); + + pool.SetInputs({input_queue_.get()}); + pool.Start(); + + // Wait for initialization to complete + pool.output_queue()->WaitForSizeAtLeast(1); + + // Now test ResetAnchor + auto [anchor, count_before] = pool.ResetAnchor(); + EXPECT_FALSE(anchor.empty()); // Should have the chunk key + EXPECT_EQ(pool.CurrentAnchor(), anchor); + EXPECT_EQ(pool.ChunksSinceAnchor(), 0); // Should be reset to 0 + + CloseInputQueue(); +} + +TEST_F(ShufflingChunkPoolTest, AnchorCounterIncrement) { + // Don't mark initial scan complete yet - we'll add sources one by one + + auto config = MakeConfig(20); + + // Start with some initial sources and complete scan + AddMockChunkSourceToQueue("source1", 20); + MarkInitialScanComplete(); + + ShufflingChunkPool pool(config); + + pool.SetInputs({input_queue_.get()}); + pool.Start(); + + // Set anchor to a key that won't match our new sources + pool.SetAnchor("non_matching_key"); + + // Wait for initial load to complete + pool.output_queue()->WaitForSizeAtLeast(1); + + // Now add new sources (these should increment the counter) + // Note: We can't add more sources after initial scan complete in the current + // setup So we'll test the counter after the initial load + + int final_count = pool.ChunksSinceAnchor(); + + // Counter should have incremented during initial load since anchor doesn't + // match + EXPECT_GT(final_count, 0); + EXPECT_EQ(pool.CurrentAnchor(), "non_matching_key"); // Anchor unchanged + + CloseInputQueue(); +} + +TEST_F(ShufflingChunkPoolTest, AnchorCounterResetDuringInitialLoad) { + // Test the special case where anchor is encountered during initial backward + // processing + AddMockChunkSourceToQueue("source_c", 10); // newest + AddMockChunkSourceToQueue("source_b", 15); // middle + AddMockChunkSourceToQueue("source_a", 20); // oldest + + auto config = MakeConfig(45); + + ShufflingChunkPool pool(config); + + pool.SetInputs({input_queue_.get()}); + pool.Start(); + + // Set anchor to middle source before marking scan complete + pool.SetAnchor("source_b"); + + // Mark scan complete to trigger initial processing + MarkInitialScanComplete(); + + // Wait for initial load to complete + pool.output_queue()->WaitForSizeAtLeast(1); + + int final_count = pool.ChunksSinceAnchor(); + + // Should only count chunks from source_c (10 chunks) since it is newer than + // the anchor. + EXPECT_EQ(final_count, 10); + EXPECT_EQ(pool.CurrentAnchor(), "source_b"); + + CloseInputQueue(); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/shuffling_frame_sampler.cc b/csrc/loader/stages/shuffling_frame_sampler.cc new file mode 100644 index 00000000..d5a75d9f --- /dev/null +++ b/csrc/loader/stages/shuffling_frame_sampler.cc @@ -0,0 +1,108 @@ +#include "loader/stages/shuffling_frame_sampler.h" + +#include + +#include "absl/algorithm/container.h" +#include "absl/log/log.h" +#include "absl/random/uniform_int_distribution.h" +#include "loader/data_loader_metrics.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" + +namespace lczero { +namespace training { + +ShufflingFrameSampler::ShufflingFrameSampler( + const ShufflingFrameSamplerConfig& config) + : SingleInputStage(config), + SingleOutputStage(config.output()), + reservoir_size_per_thread_(config.reservoir_size_per_thread()), + thread_pool_(config.threads(), ThreadPoolOptions{}) { + LOG(INFO) << "Initializing ShufflingFrameSampler with " << config.threads() + << " threads, reservoir size " + << config.reservoir_size_per_thread(); + + // Initialize thread contexts but don't start worker threads yet. + thread_contexts_.reserve(config.threads()); + for (size_t i = 0; i < config.threads(); ++i) { + thread_contexts_.push_back(std::make_unique()); + } +} + +ShufflingFrameSampler::~ShufflingFrameSampler() { Stop(); } + +void ShufflingFrameSampler::Start() { + LOG(INFO) << "Starting ShufflingFrameSampler worker threads."; + for (size_t i = 0; i < thread_contexts_.size(); ++i) { + thread_pool_.Enqueue([this, i](std::stop_token stop_token) { + Worker(stop_token, thread_contexts_[i].get()); + }); + } +} + +void ShufflingFrameSampler::Stop() { + if (thread_pool_.stop_token().stop_requested()) return; + + LOG(INFO) << "Stopping ShufflingFrameSampler."; + thread_pool_.Shutdown(); + output_queue()->Close(); + LOG(INFO) << "ShufflingFrameSampler stopped."; +} + +void ShufflingFrameSampler::Worker(std::stop_token stop_token, + ThreadContext* context) { + // Create producer early so that if input queue closes during reservoir + // prefilling, the producer will be destroyed and close the output queue. + auto producer = output_queue()->CreateProducer(); + absl::FixedArray reservoir(reservoir_size_per_thread_); + + try { + // Phase 1: Prefill the reservoir + LOG(INFO) << "ShufflingFrameSampler worker prefilling reservoir"; + absl::c_generate(reservoir, [this, context, stop_token]() { + LoadMetricPauser pauser(context->load_metric_updater); + return input_queue()->Get(stop_token); + }); + + // Phase 2: Main sampling loop + MainSamplingLoop(stop_token, reservoir, producer, context); + } catch (const QueueClosedException&) { + LOG(INFO) << "ShufflingFrameSampler worker stopping, queue closed."; + } catch (const QueueRequestCancelled&) { + LOG(INFO) << "ShufflingFrameSampler worker stopping, request cancelled."; + } +} + +void ShufflingFrameSampler::MainSamplingLoop( + std::stop_token stop_token, absl::FixedArray& reservoir, + Queue::Producer& producer, ThreadContext* context) { + absl::uniform_int_distribution dist(0, reservoir.size() - 1); + + while (true) { + const size_t random_index = dist(gen_); + { + LoadMetricPauser pauser(context->load_metric_updater); + producer.Put(std::move(reservoir[random_index]), stop_token); + } + { + LoadMetricPauser pauser(context->load_metric_updater); + reservoir[random_index] = input_queue()->Get(stop_token); + } + } +} + +StageMetricProto ShufflingFrameSampler::FlushMetrics() { + StageMetricProto stage_metric; + LoadMetricProto aggregated_load; + aggregated_load.set_name("load"); + for (const auto& context : thread_contexts_) { + UpdateFrom(aggregated_load, context->load_metric_updater.FlushMetrics()); + } + *stage_metric.add_load_metrics() = std::move(aggregated_load); + *stage_metric.add_queue_metrics() = + MetricsFromQueue("output", *output_queue()); + return stage_metric; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/shuffling_frame_sampler.h b/csrc/loader/stages/shuffling_frame_sampler.h new file mode 100644 index 00000000..a51437ec --- /dev/null +++ b/csrc/loader/stages/shuffling_frame_sampler.h @@ -0,0 +1,61 @@ +// ABOUTME: Stage that provides shuffled frames using reservoir sampling. +// ABOUTME: Takes FrameType frames and outputs them in randomized order. +#pragma once + +#include +#include +#include +#include +#include + +#include "absl/container/fixed_array.h" +#include "absl/random/random.h" +#include "loader/data_loader_metrics.h" +#include "loader/frame_type.h" +#include "loader/stages/stage.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/queue.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +// Worker that implements reservoir sampling for training frames. +// Takes FrameType frames as input and outputs them in shuffled order +// using reservoir sampling algorithm. +class ShufflingFrameSampler + : public SingleInputStage, + public SingleOutputStage { + public: + using InputType = FrameType; + using OutputType = FrameType; + + explicit ShufflingFrameSampler(const ShufflingFrameSamplerConfig& config); + ~ShufflingFrameSampler(); + + void Start() override; + void Stop() override; + StageMetricProto FlushMetrics() override; + + private: + struct ThreadContext { + LoadMetricUpdater load_metric_updater; + }; + + void Worker(std::stop_token stop_token, ThreadContext* context); + void MainSamplingLoop(std::stop_token stop_token, + absl::FixedArray& reservoir, + Queue::Producer& producer, + ThreadContext* context); + + size_t reservoir_size_per_thread_; + absl::BitGen gen_; + // thread_contexts_ must be declared before thread_pool_ to ensure + // thread_pool_ is destroyed first (stopping threads before contexts). + std::vector> thread_contexts_; + ThreadPool thread_pool_; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/shuffling_frame_sampler_test.cc b/csrc/loader/stages/shuffling_frame_sampler_test.cc new file mode 100644 index 00000000..00cd4b72 --- /dev/null +++ b/csrc/loader/stages/shuffling_frame_sampler_test.cc @@ -0,0 +1,219 @@ +#include "loader/stages/shuffling_frame_sampler.h" + +#include +#include + +#include "gtest/gtest.h" +#include "libs/lc0/src/trainingdata/trainingdata_v6.h" +#include "utils/queue.h" + +namespace lczero { +namespace training { + +namespace { + +template +class PassthroughStage : public Stage { + public: + explicit PassthroughStage(Queue* queue) : queue_(queue) {} + + void Start() override {} + void Stop() override {} + StageMetricProto FlushMetrics() override { return StageMetricProto(); } + QueueBase* GetOutput(std::string_view name = "") override { + (void)name; + return queue_; + } + void SetInputs(absl::Span inputs) override { + if (!inputs.empty()) { + throw std::runtime_error("PassthroughStage expects no inputs"); + } + } + + private: + Queue* queue_; +}; + +} // namespace + +class ShufflingFrameSamplerTest : public ::testing::Test { + protected: + void SetUp() override { + input_queue_ = std::make_unique>(100); + config_.set_reservoir_size_per_thread(10); // Small size for testing + config_.mutable_output()->set_queue_capacity(20); + } + + FrameType CreateTestFrame(uint32_t version) { + FrameType frame{}; + frame.version = version; + frame.input_format = 3; + frame.root_q = 0.5f; + return frame; + } + + std::unique_ptr> input_queue_; + ShufflingFrameSamplerConfig config_; +}; + +TEST_F(ShufflingFrameSamplerTest, OutputsNoFramesWithSmallInput) { + ShufflingFrameSampler sampler(config_); + sampler.SetInputs({input_queue_.get()}); + sampler.Start(); + + // Send 5 frames (less than reservoir size) + auto producer = input_queue_->CreateProducer(); + std::vector input_versions = {1, 2, 3, 4, 5}; + for (auto version : input_versions) { + producer.Put(CreateTestFrame(version)); + } + producer.Close(); + + // Collect all output frames + std::set output_versions; + try { + while (true) { + auto frame = sampler.output_queue()->Get(); + output_versions.insert(frame.version); + } + } catch (const QueueClosedException&) { + // Expected when queue is closed + } + + // With fewer inputs than reservoir size, no frames should be output + // (they remain in the reservoir) + EXPECT_EQ(output_versions.size(), 0); +} + +TEST_F(ShufflingFrameSamplerTest, OutputsFramesWithLargeInput) { + ShufflingFrameSampler sampler(config_); + sampler.SetInputs({input_queue_.get()}); + sampler.Start(); + + // Send 20 frames (more than reservoir size of 10) + auto producer = input_queue_->CreateProducer(); + std::vector input_versions; + for (uint32_t i = 1; i <= 20; ++i) { + input_versions.push_back(i); + producer.Put(CreateTestFrame(i)); + } + producer.Close(); + + // Collect all output frames + std::set output_versions; + try { + while (true) { + auto frame = sampler.output_queue()->Get(); + output_versions.insert(frame.version); + } + } catch (const QueueClosedException&) { + // Expected when queue is closed + } + + // Should output exactly 11 frames (10 during sampling + 1 final frame before + // queue closes) + EXPECT_EQ(output_versions.size(), 11); + + // All output frames should be from the input set + for (auto version : output_versions) { + EXPECT_TRUE(std::find(input_versions.begin(), input_versions.end(), + version) != input_versions.end()); + } +} + +TEST_F(ShufflingFrameSamplerTest, HandlesEmptyInput) { + ShufflingFrameSampler sampler(config_); + sampler.SetInputs({input_queue_.get()}); + sampler.Start(); + + // Close input queue without sending data + input_queue_->Close(); + + // Should not output any frames + EXPECT_THROW(sampler.output_queue()->Get(), QueueClosedException); +} + +TEST_F(ShufflingFrameSamplerTest, HandlesExactReservoirSize) { + ShufflingFrameSampler sampler(config_); + sampler.SetInputs({input_queue_.get()}); + sampler.Start(); + + // Send exactly reservoir_size_per_thread frames + auto producer = input_queue_->CreateProducer(); + std::vector input_versions; + for (uint32_t i = 1; i <= config_.reservoir_size_per_thread(); ++i) { + input_versions.push_back(i); + producer.Put(CreateTestFrame(i)); + } + producer.Close(); + + // Collect all output frames + std::set output_versions; + try { + while (true) { + auto frame = sampler.output_queue()->Get(); + output_versions.insert(frame.version); + } + } catch (const QueueClosedException&) { + // Expected when queue is closed + } + + // With exactly reservoir size frames, 1 frame should be output + // (fills reservoir, then queue closes during first sampling attempt) + EXPECT_EQ(output_versions.size(), 1); +} + +TEST_F(ShufflingFrameSamplerTest, PreservesFrameData) { + config_.set_reservoir_size_per_thread(2); + ShufflingFrameSampler sampler(config_); + sampler.SetInputs({input_queue_.get()}); + sampler.Start(); + + auto producer = input_queue_->CreateProducer(); + + // Create frames with specific data - need more than reservoir size + FrameType frame1 = CreateTestFrame(100); + frame1.root_q = 0.1f; + frame1.input_format = 1; + + FrameType frame2 = CreateTestFrame(200); + frame2.root_q = 0.2f; + frame2.input_format = 2; + + FrameType frame3 = CreateTestFrame(300); + frame3.root_q = 0.3f; + frame3.input_format = 3; + + producer.Put(frame1); + producer.Put(frame2); + producer.Put(frame3); // This will cause frame1 to be output + producer.Close(); + + // Verify frame data is preserved + std::vector output_frames; + try { + while (true) { + output_frames.push_back(sampler.output_queue()->Get()); + } + } catch (const QueueClosedException&) { + // Expected + } + + EXPECT_EQ(output_frames.size(), 2); + // Should be frames that were displaced from the reservoir during sampling + std::set output_frame_versions; + for (const auto& frame : output_frames) { + output_frame_versions.insert(frame.version); + // Verify frame data is preserved + if (frame.version == 100) { + EXPECT_EQ(frame.root_q, 0.1f); + EXPECT_EQ(frame.input_format, 1); + } else if (frame.version == 200) { + EXPECT_EQ(frame.root_q, 0.2f); + EXPECT_EQ(frame.input_format, 2); + } + } +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/simple_chunk_extractor.cc b/csrc/loader/stages/simple_chunk_extractor.cc new file mode 100644 index 00000000..f8e4ebbb --- /dev/null +++ b/csrc/loader/stages/simple_chunk_extractor.cc @@ -0,0 +1,106 @@ +#include "loader/stages/simple_chunk_extractor.h" + +#include +#include + +#include + +#include "loader/data_loader_metrics.h" + +namespace lczero { +namespace training { + +SimpleChunkExtractor::SimpleChunkExtractor( + const SimpleChunkExtractorConfig& config) + : SingleInputStage( + config), + SingleOutputStage(config.output()), + bitgen_(absl::MakeSeedSeq()) {} + +SimpleChunkExtractor::~SimpleChunkExtractor() { Stop(); } + +void SimpleChunkExtractor::Start() { + thread_pool_.Enqueue( + [this](std::stop_token stop_token) { Worker(stop_token); }); +} + +void SimpleChunkExtractor::Stop() { + if (thread_pool_.stop_token().stop_requested()) return; + LOG(INFO) << "Stopping SimpleChunkExtractor."; + thread_pool_.Shutdown(); + output_queue()->Close(); +} + +void SimpleChunkExtractor::Worker(std::stop_token stop_token) { + auto producer = output_queue()->CreateProducer(); + + try { + while (true) { + auto item = input_queue()->Get(stop_token); + if (item.message_type != FilePathProvider::MessageType::kFile || + !item.source) { + continue; + } + + ProcessSource(producer, std::move(item.source), stop_token); + } + } catch (const QueueClosedException&) { + } +} + +void SimpleChunkExtractor::ProcessSource( + Queue::Producer& producer, + std::unique_ptr source, std::stop_token stop_token) { + const size_t chunk_count = source->GetChunkCount(); + if (chunk_count == 0) return; + + std::vector indices(chunk_count); + std::iota(indices.begin(), indices.end(), 0); + absl::c_shuffle(indices, bitgen_); + + const std::string sort_key = source->GetChunkSortKey(); + for (size_t idx : indices) { + if (auto chunk = LoadChunk(*source, sort_key, idx)) { + producer.Put(std::move(*chunk), stop_token); + ++chunks_processed_; + } + } + ++sources_processed_; +} + +std::optional SimpleChunkExtractor::LoadChunk( + ChunkSource& source, const std::string& sort_key, size_t index) { + auto data = source.GetChunkData(index); + if (!data || data->empty()) { + ++chunks_dropped_; + return std::nullopt; + } + + TrainingChunk chunk; + chunk.sort_key = sort_key; + chunk.index_within_sort_key = index; + chunk.global_index = chunks_processed_; + chunk.use_count = 0; + chunk.frames = std::move(*data); + + return chunk; +} + +StageMetricProto SimpleChunkExtractor::FlushMetrics() { + StageMetricProto metric; + auto add_count = [&](const char* name, std::atomic& counter) { + auto* m = metric.add_count_metrics(); + m->set_name(name); + m->set_count(counter.exchange(0)); + }; + + add_count("chunks_processed", chunks_processed_); + add_count("chunks_dropped", chunks_dropped_); + add_count("sources_processed", sources_processed_); + + *metric.add_queue_metrics() = MetricsFromQueue("output", *output_queue()); + return metric; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/simple_chunk_extractor.h b/csrc/loader/stages/simple_chunk_extractor.h new file mode 100644 index 00000000..395b02ca --- /dev/null +++ b/csrc/loader/stages/simple_chunk_extractor.h @@ -0,0 +1,52 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "absl/random/random.h" +#include "loader/chunk_source/chunk_source.h" +#include "loader/stages/chunk_source_loader.h" +#include "loader/stages/stage.h" +#include "loader/stages/training_chunk.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/queue.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +// Single-threaded stage that shuffles chunks within each source. +class SimpleChunkExtractor + : public SingleInputStage, + public SingleOutputStage { + public: + explicit SimpleChunkExtractor(const SimpleChunkExtractorConfig& config); + ~SimpleChunkExtractor(); + + void Start() override; + void Stop() override; + StageMetricProto FlushMetrics() override; + + private: + void Worker(std::stop_token stop_token); + void ProcessSource(Queue::Producer& producer, + std::unique_ptr source, + std::stop_token stop_token); + std::optional LoadChunk(ChunkSource& source, + const std::string& sort_key, + size_t index); + + std::atomic chunks_processed_{0}; + std::atomic chunks_dropped_{0}; + std::atomic sources_processed_{0}; + absl::BitGen bitgen_; + ThreadPool thread_pool_{1}; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/simple_chunk_extractor_test.cc b/csrc/loader/stages/simple_chunk_extractor_test.cc new file mode 100644 index 00000000..753d1ea8 --- /dev/null +++ b/csrc/loader/stages/simple_chunk_extractor_test.cc @@ -0,0 +1,247 @@ +#include "loader/stages/simple_chunk_extractor.h" + +#include +#include + +#include +#include + +#include "loader/chunk_source/chunk_source.h" +#include "loader/stages/chunk_source_loader.h" +#include "loader/stages/file_path_provider.h" +#include "loader/stages/stage.h" +#include "loader/stages/training_chunk.h" +#include "proto/data_loader_config.pb.h" +#include "utils/queue.h" + +namespace lczero { +namespace training { +namespace { + +// Mock chunk source for testing. +class MockChunkSource : public ChunkSource { + public: + MockChunkSource(std::string sort_key, size_t chunk_count) + : sort_key_(std::move(sort_key)), chunk_count_(chunk_count) { + // Pre-generate chunk data. + for (size_t i = 0; i < chunk_count; ++i) { + chunks_.emplace_back(10); // 10 frames per chunk. + } + } + + std::string GetChunkSortKey() const override { return sort_key_; } + size_t GetChunkCount() const override { return chunk_count_; } + + std::optional> GetChunkData(size_t index) override { + if (index >= chunks_.size()) return std::nullopt; + return chunks_[index]; + } + + private: + std::string sort_key_; + size_t chunk_count_; + std::vector> chunks_; +}; + +class SimpleChunkExtractorTest : public ::testing::Test { + protected: + void SetUp() override { + // Create input queue. + input_queue_ = std::make_unique>(10); + + // Add a dummy stage to the registry. + StageConfig dummy_config; + dummy_config.set_name("dummy_input"); + registry_.AddStage("dummy_input", + std::make_unique(input_queue_.get())); + + // Create the shuffler config. + SimpleChunkExtractorConfig config; + config.set_input("dummy_input"); + config.set_queue_capacity(10); + + // Create the shuffler stage. + shuffler_ = std::make_unique(config, registry_); + } + + void TearDown() override { + if (shuffler_) { + shuffler_->Stop(); + } + input_queue_->Close(); + } + + // Helper class to provide a dummy stage for the registry. + class DummyStage : public Stage { + public: + explicit DummyStage(QueueBase* queue) : queue_(queue) {} + void Start() override {} + void Stop() override {} + StageMetricProto FlushMetrics() override { return {}; } + QueueBase* GetOutput(std::string_view name = "") override { + (void)name; + return queue_; + } + + private: + QueueBase* queue_; + }; + + StageRegistry registry_; + std::unique_ptr> input_queue_; + std::unique_ptr shuffler_; +}; + +TEST_F(SimpleChunkExtractorTest, ProcessesSingleSource) { + shuffler_->Start(); + + auto producer = input_queue_->CreateProducer(); + + // Send a chunk source with 5 chunks. + auto source = std::make_unique("source1", 5); + producer.Put({.source = std::move(source), + .message_type = FilePathProvider::MessageType::kFile}); + + // Close input to signal completion. + input_queue_->Close(); + + // Collect all output chunks. + auto* output = static_cast*>(shuffler_->GetOutput()); + std::vector chunks; + while (true) { + try { + chunks.push_back(output->Get()); + } catch (const QueueClosedException&) { + break; + } + } + + // Should receive exactly 5 chunks. + EXPECT_EQ(chunks.size(), 5); + + // All chunks should have the same sort_key. + for (const auto& chunk : chunks) { + EXPECT_EQ(chunk.sort_key, "source1"); + EXPECT_EQ(chunk.frames.size(), 10); // 10 frames per chunk. + } + + // Check that all chunk indices are present (though order is shuffled). + std::vector indices; + for (const auto& chunk : chunks) { + indices.push_back(chunk.index_within_sort_key); + } + std::sort(indices.begin(), indices.end()); + EXPECT_THAT(indices, ::testing::ElementsAre(0, 1, 2, 3, 4)); +} + +TEST_F(SimpleChunkExtractorTest, ProcessesMultipleSources) { + shuffler_->Start(); + + auto producer = input_queue_->CreateProducer(); + + // Send two chunk sources. + producer.Put({.source = std::make_unique("source1", 3), + .message_type = FilePathProvider::MessageType::kFile}); + producer.Put({.source = std::make_unique("source2", 2), + .message_type = FilePathProvider::MessageType::kFile}); + + input_queue_->Close(); + + // Collect all output chunks. + auto* output = static_cast*>(shuffler_->GetOutput()); + std::vector chunks; + while (true) { + try { + chunks.push_back(output->Get()); + } catch (const QueueClosedException&) { + break; + } + } + + // Should receive 3 + 2 = 5 chunks total. + EXPECT_EQ(chunks.size(), 5); + + // Count chunks per source. + size_t source1_count = 0; + size_t source2_count = 0; + for (const auto& chunk : chunks) { + if (chunk.sort_key == "source1") { + ++source1_count; + } else if (chunk.sort_key == "source2") { + ++source2_count; + } + } + + EXPECT_EQ(source1_count, 3); + EXPECT_EQ(source2_count, 2); +} + +TEST_F(SimpleChunkExtractorTest, SkipsNonFileMessages) { + shuffler_->Start(); + + auto producer = input_queue_->CreateProducer(); + + // Send a non-file message. + producer.Put( + {.source = nullptr, + .message_type = FilePathProvider::MessageType::kInitialScanComplete}); + + // Send a file message. + producer.Put({.source = std::make_unique("source1", 2), + .message_type = FilePathProvider::MessageType::kFile}); + + input_queue_->Close(); + + // Collect all output chunks. + auto* output = static_cast*>(shuffler_->GetOutput()); + std::vector chunks; + while (true) { + try { + chunks.push_back(output->Get()); + } catch (const QueueClosedException&) { + break; + } + } + + // Should only receive 2 chunks from the file message. + EXPECT_EQ(chunks.size(), 2); +} + +TEST_F(SimpleChunkExtractorTest, MetricsAreRecorded) { + shuffler_->Start(); + + auto producer = input_queue_->CreateProducer(); + producer.Put({.source = std::make_unique("source1", 3), + .message_type = FilePathProvider::MessageType::kFile}); + + input_queue_->Close(); + + // Wait for processing to complete. + auto* output = static_cast*>(shuffler_->GetOutput()); + while (true) { + try { + output->Get(); + } catch (const QueueClosedException&) { + break; + } + } + + // Flush metrics. + auto metrics = shuffler_->FlushMetrics(); + + EXPECT_GT(metrics.count_metrics_size(), 0); + + // Check that chunks_processed metric exists. + bool found_chunks_processed = false; + for (const auto& metric : metrics.count_metrics()) { + if (metric.name() == "chunks_processed") { + EXPECT_EQ(metric.count(), 3); + found_chunks_processed = true; + } + } + EXPECT_TRUE(found_chunks_processed); +} + +} // namespace +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/stage.cc b/csrc/loader/stages/stage.cc new file mode 100644 index 00000000..7f3e26b3 --- /dev/null +++ b/csrc/loader/stages/stage.cc @@ -0,0 +1,36 @@ +#include "loader/stages/stage.h" + +#include + +namespace lczero { +namespace training { + +void StageRegistry::AddStage(std::string_view stage_name, + std::unique_ptr stage) { + if (absl::c_find_if(stages_, [&](const auto& pair) { + return pair.first == stage_name; + }) != stages_.end()) { + throw std::runtime_error( + absl::StrCat("Duplicate stage name detected: ", stage_name)); + } + + stages_.emplace_back(stage_name, std::move(stage)); +} + +QueueBase* StageRegistry::GetStageOutput(std::string_view stage_name) const { + auto [actual_stage_name, output_name] = [&stage_name]() { + size_t dot_pos = stage_name.find('.'); + return dot_pos == std::string_view::npos + ? std::pair{stage_name, std::string_view{}} + : std::pair{stage_name.substr(0, dot_pos), + stage_name.substr(dot_pos + 1)}; + }(); + + auto it = absl::c_find_if(stages_, [&](const auto& pair) { + return pair.first == actual_stage_name; + }); + return it != stages_.end() ? it->second->GetOutput(output_name) : nullptr; +} + +} // namespace training +} // namespace lczero \ No newline at end of file diff --git a/csrc/loader/stages/stage.h b/csrc/loader/stages/stage.h new file mode 100644 index 00000000..7201d1c3 --- /dev/null +++ b/csrc/loader/stages/stage.h @@ -0,0 +1,147 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "absl/strings/str_cat.h" +#include "proto/data_loader_config.pb.h" +#include "proto/stage_control.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/queue.h" + +namespace lczero { +namespace training { + +// Base interface implemented by all loader stages. +class Stage { + public: + virtual ~Stage() = default; + + // Starts background workers owned by the stage. + virtual void Start() = 0; + + // Requests the stage to stop and join background work. + virtual void Stop() = 0; + + // Flushes stage-specific metrics and returns a snapshot. + virtual StageMetricProto FlushMetrics() = 0; + + // Returns the output queue for downstream stages. + virtual QueueBase* GetOutput(std::string_view name = "") = 0; + + // Sets the input queues for this stage. Called after construction but before + // Start(). + virtual void SetInputs(absl::Span inputs) = 0; + + // Handles control-plane messages specific to the stage. + virtual std::optional Control( + const StageControlRequest& request) { + (void)request; + return std::nullopt; + } +}; + +class StageRegistry { + public: + // Registers a new stage with the given name and takes ownership of it. + void AddStage(std::string_view stage_name, std::unique_ptr stage); + + // Returns the output queue for the specified stage. + // If stage_name contains a dot (e.g., "stage.output"), splits it into stage + // name and output name, passing the output name to Stage::GetOutput(). + // Returns nullptr if the stage is not found. + QueueBase* GetStageOutput(std::string_view stage_name) const; + + template + Queue* GetTypedStageOutput(std::string_view stage_name) const { + QueueBase* raw_queue = GetStageOutput(stage_name); + if (raw_queue == nullptr) return nullptr; + auto* typed_queue = dynamic_cast*>(raw_queue); + if (!typed_queue) { + throw std::runtime_error( + absl::StrCat("Stage output type mismatch for stage: ", stage_name)); + } + return typed_queue; + } + + size_t size() const { return stages_.size(); } + const std::vector>>& stages() + const { + return stages_; + } + + private: + std::vector>> stages_; +}; + +// Helper to convert QueueConfig::OverflowBehavior to OverflowBehavior. +inline OverflowBehavior ToOverflowBehavior( + QueueConfig::OverflowBehavior behavior) { + switch (behavior) { + case QueueConfig::BLOCK: + return OverflowBehavior::BLOCK; + case QueueConfig::DROP_NEW: + return OverflowBehavior::DROP_NEW; + case QueueConfig::KEEP_NEWEST: + return OverflowBehavior::KEEP_NEWEST; + } + throw std::runtime_error(absl::StrCat("Unknown OverflowBehavior value: ", + static_cast(behavior))); +} + +// Helper for stages that consume a single upstream queue. +template +class SingleInputStage : virtual public Stage { + public: + void SetInputs(absl::Span inputs) override { + if (inputs.size() != 1) { + throw std::runtime_error(absl::StrCat( + "SingleInputStage expects exactly 1 input, got ", inputs.size())); + } + auto* typed_queue = dynamic_cast*>(inputs[0]); + if (!typed_queue) throw std::runtime_error("Input queue type mismatch"); + input_queue_ = typed_queue; + } + + protected: + explicit SingleInputStage(const ConfigT&) : input_queue_(nullptr) {} + + Queue* input_queue() { return input_queue_; } + + private: + Queue* input_queue_; +}; + +// Helper for stages that produce a single output queue. +template +class SingleOutputStage : virtual public Stage { + public: + Queue* output_queue() { return &output_queue_; } + + QueueBase* GetOutput(std::string_view name = "") override { + if (name != output_name_) { + throw std::runtime_error(absl::StrCat("Output name '", name, + "' does not match configured '", + output_name_, "'")); + } + return &output_queue_; + } + + protected: + explicit SingleOutputStage(const QueueConfig& config) + : output_name_(config.name()), + output_queue_(config.queue_capacity(), + ToOverflowBehavior(config.overflow_behavior())) {} + + private: + std::string output_name_; + Queue output_queue_; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/stage_factory.cc b/csrc/loader/stages/stage_factory.cc new file mode 100644 index 00000000..6d2617ee --- /dev/null +++ b/csrc/loader/stages/stage_factory.cc @@ -0,0 +1,82 @@ +#include "loader/stages/stage_factory.h" + +#include +#include + +#include "loader/stages/chunk_rescorer.h" +#include "loader/stages/chunk_source_loader.h" +#include "loader/stages/chunk_source_splitter.h" +#include "loader/stages/chunk_unpacker.h" +#include "loader/stages/file_path_provider.h" +#include "loader/stages/join_stage.h" +#include "loader/stages/shuffling_chunk_pool.h" +#include "loader/stages/shuffling_frame_sampler.h" +#include "loader/stages/simple_chunk_extractor.h" +#include "loader/stages/tensor_generator.h" + +namespace lczero { +namespace training { + +namespace { + +int CountStageConfigs(const StageConfig& config) { + return static_cast(config.has_file_path_provider()) + + static_cast(config.has_chunk_source_loader()) + + static_cast(config.has_shuffling_chunk_pool()) + + static_cast(config.has_chunk_rescorer()) + + static_cast(config.has_chunk_unpacker()) + + static_cast(config.has_shuffling_frame_sampler()) + + static_cast(config.has_tensor_generator()) + + static_cast(config.has_chunk_source_splitter()) + + static_cast(config.has_simple_chunk_extractor()) + + static_cast(config.has_join_positions()); +} + +} // namespace + +std::unique_ptr CreateStage(const StageConfig& config) { + if (CountStageConfigs(config) != 1) { + throw std::runtime_error( + "StageConfig must have exactly one stage-specific config set."); + } + + if (config.has_file_path_provider()) { + return std::make_unique(config.file_path_provider()); + } + if (config.has_chunk_source_loader()) { + return std::make_unique(config.chunk_source_loader()); + } + if (config.has_shuffling_chunk_pool()) { + return std::make_unique(config.shuffling_chunk_pool()); + } + if (config.has_chunk_rescorer()) { + return std::make_unique(config.chunk_rescorer()); + } + if (config.has_chunk_unpacker()) { + return std::make_unique(config.chunk_unpacker()); + } + if (config.has_shuffling_frame_sampler()) { + return std::make_unique( + config.shuffling_frame_sampler()); + } + if (config.has_tensor_generator()) { + return std::make_unique(config.tensor_generator()); + } + if (config.has_chunk_source_splitter()) { + return std::make_unique( + config.chunk_source_splitter()); + } + if (config.has_simple_chunk_extractor()) { + return std::make_unique( + config.simple_chunk_extractor()); + } + if (config.has_join_positions()) { + return std::make_unique(config.join_positions()); + } + + throw std::runtime_error( + "StageConfig did not contain a recognized stage configuration."); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/stage_factory.h b/csrc/loader/stages/stage_factory.h new file mode 100644 index 00000000..f74acd6b --- /dev/null +++ b/csrc/loader/stages/stage_factory.h @@ -0,0 +1,14 @@ +#pragma once + +#include + +#include "loader/stages/stage.h" +#include "proto/data_loader_config.pb.h" + +namespace lczero { +namespace training { + +std::unique_ptr CreateStage(const StageConfig& config); + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/stage_factory_test.cc b/csrc/loader/stages/stage_factory_test.cc new file mode 100644 index 00000000..a0c5b042 --- /dev/null +++ b/csrc/loader/stages/stage_factory_test.cc @@ -0,0 +1,35 @@ +#include "loader/stages/stage_factory.h" + +#include + +#include + +namespace lczero { +namespace training { + +TEST(StageFactoryTest, CreatesFilePathProviderStage) { + StageConfig config; + config.mutable_file_path_provider()->set_directory("."); + + auto stage = CreateStage(config); + + ASSERT_NE(stage, nullptr); + EXPECT_NE(stage->GetOutput(), nullptr); +} + +TEST(StageFactoryTest, ThrowsWhenNoStageConfigSet) { + StageConfig config; + + EXPECT_THROW(CreateStage(config), std::runtime_error); +} + +TEST(StageFactoryTest, ThrowsWhenMultipleStageConfigsSet) { + StageConfig config; + config.mutable_file_path_provider()->set_directory("."); + config.mutable_tensor_generator(); + + EXPECT_THROW(CreateStage(config), std::runtime_error); +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/tensor_generator.cc b/csrc/loader/stages/tensor_generator.cc new file mode 100644 index 00000000..145bad86 --- /dev/null +++ b/csrc/loader/stages/tensor_generator.cc @@ -0,0 +1,216 @@ +// ABOUTME: Implementation of TensorGenerator stage for training pipeline. +// ABOUTME: Converts V6TrainingData frames to batched tensors for training. + +#include "loader/stages/tensor_generator.h" + +#include +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/log/log.h" +#include "loader/data_loader_metrics.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" + +namespace lczero { +namespace training { + +TensorGenerator::TensorGenerator(const TensorGeneratorConfig& config) + : SingleInputStage(config), + SingleOutputStage(config.output()), + batch_size_(config.batch_size()), + thread_pool_(config.threads(), ThreadPoolOptions{}) { + LOG(INFO) << "Initializing TensorGenerator with " << config.threads() + << " threads, batch size " << config.batch_size(); + + // Initialize thread contexts but don't start worker threads yet. + thread_contexts_.reserve(config.threads()); + for (size_t i = 0; i < config.threads(); ++i) { + thread_contexts_.push_back(std::make_unique()); + } +} + +TensorGenerator::~TensorGenerator() { Stop(); } + +void TensorGenerator::Start() { + LOG(INFO) << "Starting TensorGenerator worker threads."; + for (size_t i = 0; i < thread_contexts_.size(); ++i) { + thread_pool_.Enqueue([this, i](std::stop_token stop_token) { + Worker(stop_token, thread_contexts_[i].get()); + }); + } +} + +void TensorGenerator::Stop() { + if (thread_pool_.stop_token().stop_requested()) return; + + LOG(INFO) << "Stopping TensorGenerator."; + thread_pool_.Shutdown(); + output_queue()->Close(); + LOG(INFO) << "TensorGenerator stopped."; +} + +void TensorGenerator::Worker(std::stop_token stop_token, + ThreadContext* context) { + auto producer = output_queue()->CreateProducer(); + std::vector batch; + batch.reserve(batch_size_); + + try { + while (true) { + // Collect frames for a batch. + batch.clear(); + for (size_t i = 0; i < batch_size_; ++i) { + LoadMetricPauser pauser(context->load_metric_updater); + batch.push_back(input_queue()->Get(stop_token)); + } + + // Convert batch to tensors. + TensorTuple tensors = ConvertFramesToTensors(batch); + { + LoadMetricPauser pauser(context->load_metric_updater); + producer.Put(std::move(tensors), stop_token); + } + } + } catch (const QueueClosedException&) { + LOG(INFO) << "TensorGenerator worker stopping, queue closed."; + } catch (const QueueRequestCancelled&) { + LOG(INFO) << "TensorGenerator worker stopping, request cancelled."; + } +} + +TensorTuple TensorGenerator::ConvertFramesToTensors( + const std::vector& frames) { + const size_t batch_size = frames.size(); + constexpr size_t kNumPlanes = 112; + constexpr size_t kNumPolicyMoves = 1858; + constexpr size_t kNumValueTypes = 6; + constexpr size_t kValuesPerType = 3; + + TensorTuple result; + result.reserve(3); + + // Index 0: Input planes (batch_size, 112, 8, 8) + auto planes_tensor = std::make_unique>( + std::initializer_list{batch_size, kNumPlanes, 8, 8}); + ProcessPlanes(frames, *planes_tensor); + result.push_back(std::move(planes_tensor)); + + // Index 1: Probabilities (batch_size, 1858) + auto probs_tensor = std::make_unique>( + std::initializer_list{batch_size, kNumPolicyMoves}); + for (size_t i = 0; i < batch_size; ++i) { + auto probs_slice = probs_tensor->slice({static_cast(i)}); + std::memcpy(probs_slice.data(), frames[i].probabilities, + kNumPolicyMoves * sizeof(float)); + } + result.push_back(std::move(probs_tensor)); + + // Index 2: Values (batch_size, 6, 3) with [q, d, m] for each type. + // [0]: result, [1]: best, [2]: played, [3]: orig, [4]: root, [5]: st + auto values_tensor = + std::make_unique>(std::initializer_list{ + batch_size, kNumValueTypes, kValuesPerType}); + for (size_t i = 0; i < batch_size; ++i) { + const auto& frame = frames[i]; + auto batch_slice = values_tensor->slice({static_cast(i)}); + + // Index 0: result [result_q, result_d, plies_left] + auto result_slice = batch_slice.subspan(0 * kValuesPerType, kValuesPerType); + result_slice[0] = frame.result_q; + result_slice[1] = frame.result_d; + result_slice[2] = frame.plies_left; + + // Index 1: best [best_q, best_d, best_m] + auto best_slice = batch_slice.subspan(1 * kValuesPerType, kValuesPerType); + best_slice[0] = frame.best_q; + best_slice[1] = frame.best_d; + best_slice[2] = frame.best_m; + + // Index 2: played [played_q, played_d, played_m] + auto played_slice = batch_slice.subspan(2 * kValuesPerType, kValuesPerType); + played_slice[0] = frame.played_q; + played_slice[1] = frame.played_d; + played_slice[2] = frame.played_m; + + // Index 3: orig [orig_q, orig_d, orig_m] (may be NaN) + auto orig_slice = batch_slice.subspan(3 * kValuesPerType, kValuesPerType); + orig_slice[0] = frame.orig_q; + orig_slice[1] = frame.orig_d; + orig_slice[2] = frame.orig_m; + + // Index 4: root [root_q, root_d, root_m] + auto root_slice = batch_slice.subspan(4 * kValuesPerType, kValuesPerType); + root_slice[0] = frame.root_q; + root_slice[1] = frame.root_d; + root_slice[2] = frame.root_m; + + // Index 5: st [q_st, d_st, NaN] + auto st_slice = batch_slice.subspan(5 * kValuesPerType, kValuesPerType); + st_slice[0] = frame.q_st; + st_slice[1] = frame.d_st; + st_slice[2] = std::numeric_limits::quiet_NaN(); + } + result.push_back(std::move(values_tensor)); + + return result; +} + +void TensorGenerator::ProcessPlanes(const std::vector& frames, + TypedTensor& planes_tensor) { + const size_t batch_size = frames.size(); + + for (size_t i = 0; i < batch_size; ++i) { + const auto& frame = frames[i]; + auto batch_slice = planes_tensor.slice({static_cast(i)}); + + // Process first 104 planes from frame.planes (each uint64_t represents 64 + // bits). + for (ssize_t plane = 0; plane < 104; ++plane) { + auto plane_slice = batch_slice.subspan(plane * 64, 64); + uint64_t plane_bits = frame.planes[plane]; + + for (ssize_t square = 0; square < 64; ++square) { + // XOR with 7 remaps the index within each byte from 0..7 to 7..0. + plane_slice[square] = + static_cast((plane_bits >> (square ^ 7)) & 1); + } + } + + // Add 8 additional planes for metadata (planes 104-111). + const std::pair meta_planes[] = { + {104, static_cast(frame.castling_us_ooo)}, + {105, static_cast(frame.castling_us_oo)}, + {106, static_cast(frame.castling_them_ooo)}, + {107, static_cast(frame.castling_them_oo)}, + {108, static_cast(frame.side_to_move_or_enpassant)}, + {109, static_cast(frame.rule50_count) / 99.0f}, + {110, 0.0f}, // All zeros (constant plane). + {111, 1.0f}, // All ones (constant plane). + }; + + for (const auto& [plane_num, value] : meta_planes) { + auto plane_slice = batch_slice.subspan(plane_num * 64, 64); + absl::c_fill(plane_slice, value); + } + } +} + +StageMetricProto TensorGenerator::FlushMetrics() { + StageMetricProto stage_metric; + LoadMetricProto aggregated_load; + aggregated_load.set_name("load"); + for (const auto& context : thread_contexts_) { + UpdateFrom(aggregated_load, context->load_metric_updater.FlushMetrics()); + } + *stage_metric.add_load_metrics() = std::move(aggregated_load); + *stage_metric.add_queue_metrics() = + MetricsFromQueue("output", *output_queue()); + return stage_metric; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/tensor_generator.h b/csrc/loader/stages/tensor_generator.h new file mode 100644 index 00000000..c1ba3033 --- /dev/null +++ b/csrc/loader/stages/tensor_generator.h @@ -0,0 +1,59 @@ +// ABOUTME: Stage that converts FrameType frames into tensor batches. +// ABOUTME: Produces TrainingTensors with tensors for training pipeline. +#pragma once + +#include +#include +#include +#include +#include + +#include "loader/data_loader.h" +#include "loader/data_loader_metrics.h" +#include "loader/frame_type.h" +#include "loader/stages/stage.h" +#include "proto/data_loader_config.pb.h" +#include "proto/training_metrics.pb.h" +#include "utils/queue.h" +#include "utils/tensor.h" +#include "utils/thread_pool.h" + +namespace lczero { +namespace training { + +// Worker pool that converts FrameType frames into tensor batches. +// Takes individual FrameType frames as input and outputs TensorTuple +// containing batched tensors in the format required for training. +class TensorGenerator + : public SingleInputStage, + public SingleOutputStage { + public: + using InputType = FrameType; + using OutputType = TensorTuple; + + explicit TensorGenerator(const TensorGeneratorConfig& config); + ~TensorGenerator(); + + void Start() override; + void Stop() override; + StageMetricProto FlushMetrics() override; + + private: + struct ThreadContext { + LoadMetricUpdater load_metric_updater; + }; + + void Worker(std::stop_token stop_token, ThreadContext* context); + TensorTuple ConvertFramesToTensors(const std::vector& frames); + void ProcessPlanes(const std::vector& frames, + TypedTensor& planes_tensor); + + size_t batch_size_; + // thread_contexts_ must be declared before thread_pool_ to ensure + // thread_pool_ is destroyed first (stopping threads before contexts). + std::vector> thread_contexts_; + ThreadPool thread_pool_; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/tensor_generator_test.cc b/csrc/loader/stages/tensor_generator_test.cc new file mode 100644 index 00000000..4e8f5c5a --- /dev/null +++ b/csrc/loader/stages/tensor_generator_test.cc @@ -0,0 +1,373 @@ +// ABOUTME: Unit tests for TensorGenerator stage in training pipeline. +// ABOUTME: Tests tensor conversion, batching, and data format correctness. + +#include "loader/stages/tensor_generator.h" + +#include +#include +#include + +#include "gtest/gtest.h" +#include "libs/lc0/src/trainingdata/trainingdata_v6.h" +#include "utils/queue.h" +#include "utils/tensor.h" + +namespace lczero { +namespace training { + +namespace { + +template +class PassthroughStage : public Stage { + public: + explicit PassthroughStage(Queue* queue) : queue_(queue) {} + + void Start() override {} + void Stop() override {} + StageMetricProto FlushMetrics() override { return StageMetricProto(); } + QueueBase* GetOutput(std::string_view name = "") override { + (void)name; + return queue_; + } + void SetInputs(absl::Span inputs) override { + if (!inputs.empty()) { + throw std::runtime_error("PassthroughStage expects no inputs"); + } + } + + private: + Queue* queue_; +}; + +} // namespace + +class TensorGeneratorTest : public ::testing::Test { + protected: + void SetUp() override { + input_queue_ = std::make_unique>(100); + config_.set_batch_size(4); + config_.set_threads(1); + config_.mutable_output()->set_queue_capacity(10); + } + + FrameType CreateTestFrame() { + FrameType frame{}; + std::memset(&frame, 0, sizeof(frame)); + + frame.version = 6; + frame.input_format = 3; + + // Fill probabilities with test values. + for (ssize_t i = 0; i < 1858; ++i) { + frame.probabilities[i] = static_cast(i) / 1858.0f; + } + + // Fill planes with test pattern. + for (ssize_t i = 0; i < 104; ++i) { + frame.planes[i] = 0x0F0F0F0F0F0F0F0FULL + i; // Test pattern + } + + // Set castling rights. + frame.castling_us_ooo = 1; + frame.castling_us_oo = 0; + frame.castling_them_ooo = 1; + frame.castling_them_oo = 1; + + // Set other fields. + frame.side_to_move_or_enpassant = 1; + frame.rule50_count = 50; + + // Set Q and D values. + frame.result_q = 0.5f; + frame.result_d = 0.2f; + frame.best_q = 0.3f; + frame.best_d = 0.1f; + frame.best_m = 42.5f; + frame.plies_left = 42.5f; + + return frame; + } + + void VerifyTensorTuple(const TensorTuple& tensors, + const std::vector& frames) { + const size_t batch_size = frames.size(); + + // Verify tuple has 3 elements + ASSERT_EQ(tensors.size(), 3); + + // Verify input tensor: (batch_size, 112, 8, 8) + const auto* planes_tensor = + dynamic_cast*>(tensors[0].get()); + ASSERT_NE(planes_tensor, nullptr); + EXPECT_EQ(planes_tensor->shape().size(), 4); + EXPECT_EQ(planes_tensor->shape()[0], batch_size); + EXPECT_EQ(planes_tensor->shape()[1], 112); + EXPECT_EQ(planes_tensor->shape()[2], 8); + EXPECT_EQ(planes_tensor->shape()[3], 8); + + // Verify probabilities tensor: (batch_size, 1858) + const auto* probs_tensor = + dynamic_cast*>(tensors[1].get()); + ASSERT_NE(probs_tensor, nullptr); + EXPECT_EQ(probs_tensor->shape().size(), 2); + EXPECT_EQ(probs_tensor->shape()[0], batch_size); + EXPECT_EQ(probs_tensor->shape()[1], 1858); + + // Verify values tensor: (batch_size, 6, 3) + const auto* values_tensor = + dynamic_cast*>(tensors[2].get()); + ASSERT_NE(values_tensor, nullptr); + EXPECT_EQ(values_tensor->shape().size(), 3); + EXPECT_EQ(values_tensor->shape()[0], batch_size); + EXPECT_EQ(values_tensor->shape()[1], 6); + EXPECT_EQ(values_tensor->shape()[2], 3); + } + + void VerifyTensorData(const TensorTuple& tensors, + const std::vector& frames) { + const size_t batch_size = frames.size(); + const auto* planes_tensor = + dynamic_cast*>(tensors[0].get()); + const auto* probs_tensor = + dynamic_cast*>(tensors[1].get()); + const auto* values_tensor = + dynamic_cast*>(tensors[2].get()); + + for (size_t i = 0; i < batch_size; ++i) { + const auto& frame = frames[i]; + + // Verify probabilities data. + auto probs_slice = probs_tensor->slice({static_cast(i)}); + for (ssize_t j = 0; j < 1858; ++j) { + EXPECT_FLOAT_EQ(probs_slice[j], frame.probabilities[j]); + } + + // Verify values tensor [batch, 6, 3] with raw q/d/m values + // Index 0: result (q=0.5, d=0.2, m=42.5) + auto values_slice = values_tensor->slice({static_cast(i)}); + EXPECT_FLOAT_EQ(values_slice[0 * 3 + 0], 0.5f); // result_q + EXPECT_FLOAT_EQ(values_slice[0 * 3 + 1], 0.2f); // result_d + EXPECT_FLOAT_EQ(values_slice[0 * 3 + 2], 42.5f); // result_m + + // Index 1: best (q=0.3, d=0.1, m=42.5) + EXPECT_FLOAT_EQ(values_slice[1 * 3 + 0], 0.3f); // best_q + EXPECT_FLOAT_EQ(values_slice[1 * 3 + 1], 0.1f); // best_d + EXPECT_FLOAT_EQ(values_slice[1 * 3 + 2], 42.5f); // best_m + + // Verify planes data - check first few planes and meta planes. + auto planes_slice = planes_tensor->slice({static_cast(i)}); + + // Check first plane (plane 0). + uint64_t expected_plane_0 = 0x0F0F0F0F0F0F0F0FULL; + for (ssize_t square = 0; square < 64; ++square) { + float expected = + static_cast((expected_plane_0 >> (63 - square)) & 1); + EXPECT_FLOAT_EQ(planes_slice[square], expected); + } + + // Check meta planes. + // Plane 104: castling_us_ooo = 1 + for (ssize_t square = 104 * 64; square < 105 * 64; ++square) { + EXPECT_FLOAT_EQ(planes_slice[square], 1.0f); + } + + // Plane 105: castling_us_oo = 0 + for (ssize_t square = 105 * 64; square < 106 * 64; ++square) { + EXPECT_FLOAT_EQ(planes_slice[square], 0.0f); + } + + // Plane 109: rule50_count = 50, should be 50/99 + for (ssize_t square = 109 * 64; square < 110 * 64; ++square) { + EXPECT_FLOAT_EQ(planes_slice[square], 50.0f / 99.0f); + } + + // Plane 110: all zeros + for (ssize_t square = 110 * 64; square < 111 * 64; ++square) { + EXPECT_FLOAT_EQ(planes_slice[square], 0.0f); + } + + // Plane 111: all ones + for (ssize_t square = 111 * 64; square < 112 * 64; ++square) { + EXPECT_FLOAT_EQ(planes_slice[square], 1.0f); + } + } + } + + std::unique_ptr> input_queue_; + TensorGeneratorConfig config_; +}; + +TEST_F(TensorGeneratorTest, GeneratesCorrectTensorShapes) { + TensorGenerator generator(config_); + generator.SetInputs({input_queue_.get()}); + generator.Start(); + + auto producer = input_queue_->CreateProducer(); + std::vector frames; + for (size_t i = 0; i < config_.batch_size(); ++i) { + frames.push_back(CreateTestFrame()); + producer.Put(frames.back()); + } + producer.Close(); + + auto tensors = generator.output_queue()->Get(); + VerifyTensorTuple(tensors, frames); +} + +TEST_F(TensorGeneratorTest, GeneratesCorrectTensorData) { + TensorGenerator generator(config_); + generator.SetInputs({input_queue_.get()}); + generator.Start(); + + auto producer = input_queue_->CreateProducer(); + std::vector frames; + for (size_t i = 0; i < config_.batch_size(); ++i) { + frames.push_back(CreateTestFrame()); + producer.Put(frames.back()); + } + producer.Close(); + + auto tensors = generator.output_queue()->Get(); + VerifyTensorTuple(tensors, frames); + VerifyTensorData(tensors, frames); +} + +TEST_F(TensorGeneratorTest, HandlesMultipleBatches) { + TensorGenerator generator(config_); + generator.SetInputs({input_queue_.get()}); + generator.Start(); + + auto producer = input_queue_->CreateProducer(); + + // Send two full batches. + std::vector all_frames; + for (ssize_t batch = 0; batch < 2; ++batch) { + for (size_t i = 0; i < config_.batch_size(); ++i) { + auto frame = CreateTestFrame(); + frame.version = batch * 1000 + i; // Unique version for each frame + all_frames.push_back(frame); + producer.Put(frame); + } + } + producer.Close(); + + // Get first batch. + auto tensors1 = generator.output_queue()->Get(); + std::vector batch1_frames( + all_frames.begin(), all_frames.begin() + config_.batch_size()); + VerifyTensorTuple(tensors1, batch1_frames); + + // Get second batch. + auto tensors2 = generator.output_queue()->Get(); + std::vector batch2_frames( + all_frames.begin() + config_.batch_size(), all_frames.end()); + VerifyTensorTuple(tensors2, batch2_frames); + + // No more batches should be available. + EXPECT_THROW(generator.output_queue()->Get(), QueueClosedException); +} + +TEST_F(TensorGeneratorTest, HandlesDifferentBatchSizes) { + config_.set_batch_size(2); + TensorGenerator generator(config_); + generator.SetInputs({input_queue_.get()}); + generator.Start(); + + auto producer = input_queue_->CreateProducer(); + std::vector frames; + for (size_t i = 0; i < config_.batch_size(); ++i) { + frames.push_back(CreateTestFrame()); + producer.Put(frames.back()); + } + producer.Close(); + + auto tensors = generator.output_queue()->Get(); + VerifyTensorTuple(tensors, frames); +} + +TEST_F(TensorGeneratorTest, HandlesEmptyInput) { + TensorGenerator generator(config_); + generator.SetInputs({input_queue_.get()}); + generator.Start(); + + // Close input queue without sending data. + input_queue_->Close(); + + // Should not output any tensors. + EXPECT_THROW(generator.output_queue()->Get(), QueueClosedException); +} + +TEST_F(TensorGeneratorTest, VerifiesPlanesConversion) { + config_.set_batch_size(1); + TensorGenerator generator(config_); + generator.SetInputs({input_queue_.get()}); + generator.Start(); + + auto producer = input_queue_->CreateProducer(); + + FrameType frame = CreateTestFrame(); + // Set specific bit pattern for plane 0. + frame.planes[0] = 0xAAAAAAAAAAAAAAAAULL; // Alternating bits + // Set specific values for meta planes. + frame.castling_us_ooo = 1; + frame.castling_us_oo = 0; + frame.rule50_count = 75; + + producer.Put(frame); + producer.Close(); + + auto tensors = generator.output_queue()->Get(); + const auto* planes_tensor = + dynamic_cast*>(tensors[0].get()); + + auto planes_slice = planes_tensor->slice({0}); + + // Verify plane 0 bit conversion. + for (ssize_t square = 0; square < 64; ++square) { + float expected = + static_cast((0xAAAAAAAAAAAAAAAAULL >> (63 - square)) & 1); + EXPECT_FLOAT_EQ(planes_slice[square], expected) + << "Mismatch at square " << square; + } + + // Verify rule50_count conversion: 75/99. + for (ssize_t square = 109 * 64; square < 110 * 64; ++square) { + EXPECT_FLOAT_EQ(planes_slice[square], 75.0f / 99.0f); + } +} + +TEST_F(TensorGeneratorTest, VerifiesQDConversion) { + config_.set_batch_size(1); + TensorGenerator generator(config_); + generator.SetInputs({input_queue_.get()}); + generator.Start(); + + auto producer = input_queue_->CreateProducer(); + + FrameType frame = CreateTestFrame(); + // Test specific Q/D values. + frame.result_q = 0.4f; + frame.result_d = 0.3f; + frame.best_q = -0.2f; + frame.best_d = 0.1f; + + producer.Put(frame); + producer.Close(); + + auto tensors = generator.output_queue()->Get(); + const auto* values_tensor = + dynamic_cast*>(tensors[2].get()); + + auto values_slice = values_tensor->slice({0}); + + // Verify result values: q=0.4, d=0.3 (raw values, no WDL conversion) + EXPECT_FLOAT_EQ(values_slice[0 * 3 + 0], 0.4f); // result_q + EXPECT_FLOAT_EQ(values_slice[0 * 3 + 1], 0.3f); // result_d + + // Verify best values: q=-0.2, d=0.1 (raw values, no WDL conversion) + EXPECT_FLOAT_EQ(values_slice[1 * 3 + 0], -0.2f); // best_q + EXPECT_FLOAT_EQ(values_slice[1 * 3 + 1], 0.1f); // best_d +} + +} // namespace training +} // namespace lczero diff --git a/csrc/loader/stages/training_chunk.h b/csrc/loader/stages/training_chunk.h new file mode 100644 index 00000000..dca3765a --- /dev/null +++ b/csrc/loader/stages/training_chunk.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include +#include +#include + +#include "loader/frame_type.h" + +namespace lczero { +namespace training { + +struct TrainingChunk { + std::vector frames; + std::string sort_key; + size_t index_within_sort_key = 0; + size_t global_index = 0; + uint32_t use_count = 0; +}; + +struct CacheRequest { + size_t global_index = 0; + uint16_t next_use = 0; + std::vector items; +}; + +} // namespace training +} // namespace lczero diff --git a/csrc/tools/dump_chunk_main.cc b/csrc/tools/dump_chunk_main.cc new file mode 100644 index 00000000..cac3ae7a --- /dev/null +++ b/csrc/tools/dump_chunk_main.cc @@ -0,0 +1,97 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "trainingdata/trainingdata_v6.h" +#include "utils/training_data_printer.h" + +ABSL_FLAG(std::string, chunk_path, "", "Path to the chunk file (.gz) to dump."); +ABSL_FLAG(int64_t, max_entries, -1, + "Maximum number of entries to print. -1 prints all entries."); +ABSL_FLAG(int64_t, float_values_per_line, 8, + "Number of floating point values per output line."); +ABSL_FLAG(int64_t, plane_values_per_line, 4, + "Number of plane values per output line."); + +namespace lczero { +namespace training { + +namespace { + +using ::lczero::training::FrameType; +using ::lczero::training::PrintTrainingDataEntry; + +void DumpChunk(const std::string& path, int64_t max_entries, + int64_t float_per_line, int64_t plane_per_line) { + gzFile file = gzopen(path.c_str(), "rb"); + if (file == nullptr) { + LOG(FATAL) << "Failed to open chunk file: " << path; + } + + size_t index = 0; + while (true) { + FrameType entry; + const int bytes_read = gzread(file, &entry, sizeof(entry)); + if (bytes_read == 0) { + break; + } + if (bytes_read < 0) { + int errnum = 0; + const char* error_message = gzerror(file, &errnum); + gzclose(file); + LOG(FATAL) << "Error while reading chunk: " << error_message; + } + if (bytes_read != sizeof(entry)) { + gzclose(file); + LOG(FATAL) << "Unexpected chunk size. Expected " << sizeof(entry) + << " bytes, got " << bytes_read << "."; + } + + const std::string header = absl::StrFormat("Entry %zu:", index); + PrintTrainingDataEntry(entry, header, float_per_line, plane_per_line); + ++index; + + if (max_entries >= 0 && static_cast(index) >= max_entries) { + break; + } + } + + gzclose(file); + LOG(INFO) << "Printed " << index << " entries."; +} + +} // namespace + +} // namespace training +} // namespace lczero + +int main(int argc, char** argv) { + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo); + + const std::string chunk_path = absl::GetFlag(FLAGS_chunk_path); + if (chunk_path.empty()) { + LOG(FATAL) << "--chunk_path flag is required."; + } + + const int64_t max_entries = absl::GetFlag(FLAGS_max_entries); + const int64_t float_per_line = absl::GetFlag(FLAGS_float_values_per_line); + const int64_t plane_per_line = absl::GetFlag(FLAGS_plane_values_per_line); + + lczero::training::DumpChunk(chunk_path, max_entries, float_per_line, + plane_per_line); + return 0; +} diff --git a/csrc/tools/filter_chunks_main.cc b/csrc/tools/filter_chunks_main.cc new file mode 100644 index 00000000..fbaee280 --- /dev/null +++ b/csrc/tools/filter_chunks_main.cc @@ -0,0 +1,210 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "loader/chunk_source/chunk_source.h" +#include "loader/chunk_source/tar_chunk_source.h" +#include "trainingdata/trainingdata_v6.h" + +ABSL_FLAG(std::string, input_dir, ".", "Directory to scan for .tar files."); +ABSL_FLAG(std::string, output_dir, ".", + "Directory where matching chunks will be written."); +ABSL_FLAG(std::string, plane_values, "", + "Comma separated list of plane values (decimal or hex)."); + +namespace { + +namespace fs = std::filesystem; + +using ::lczero::training::ChunkSource; +using ::lczero::training::ChunkSourceLoaderConfig; +using ::lczero::training::FrameType; +using ::lczero::training::TarChunkSource; + +std::vector CollectTarFiles(const fs::path& directory) { + std::vector files; + for (const auto& entry : fs::directory_iterator(directory)) { + const fs::path& path = entry.path(); + if (entry.is_regular_file() && path.extension() == ".tar") { + files.push_back(path); + } + } + absl::c_sort(files, [](const fs::path& lhs, const fs::path& rhs) { + return lhs.filename() < rhs.filename(); + }); + return files; +} + +std::vector ParsePlaneValues(absl::string_view value_list) { + std::vector result; + if (value_list.empty()) { + LOG(FATAL) << "--plane_values flag must not be empty."; + } + for (absl::string_view token : + absl::StrSplit(value_list, ',', absl::SkipWhitespace())) { + token = absl::StripAsciiWhitespace(token); + if (token.empty()) continue; + + uint64_t value = 0; + if (absl::StartsWithIgnoreCase(token, "0x")) { + const absl::string_view hex_part = token.substr(2); + if (hex_part.empty() || !absl::SimpleHexAtoi(hex_part, &value)) { + LOG(FATAL) << "Invalid hex plane value: " << token; + } + } else if (!absl::SimpleAtoi(token, &value)) { + LOG(FATAL) << "Invalid decimal plane value: " << token; + } + result.push_back(value); + } + if (result.empty()) { + LOG(FATAL) << "No plane values were parsed."; + } + return result; +} + +bool PlanesMatch(const FrameType& entry, absl::Span expected) { + if (expected.size() > std::size(entry.planes)) return false; + const size_t bytes = expected.size() * sizeof(uint64_t); + return std::memcmp(entry.planes, expected.data(), bytes) == 0; +} + +std::optional FindMatchingFrameIndex( + const std::vector& chunk, absl::Span expected) { + for (size_t frame = 0; frame < chunk.size(); ++frame) { + if (PlanesMatch(chunk[frame], expected)) return frame; + } + return std::nullopt; +} + +void WriteChunk(const fs::path& output_dir, absl::string_view base_name, + size_t index, size_t frame_index, + const std::vector& chunk) { + fs::create_directories(output_dir); + const fs::path output_path = + output_dir / absl::StrCat(base_name, "_", index, "_", frame_index, ".gz"); + + gzFile file = gzopen(output_path.string().c_str(), "wb"); + if (file == nullptr) { + LOG(FATAL) << "Failed to open output file: " << output_path.string(); + } + + size_t remaining = chunk.size() * sizeof(FrameType); + const char* data = reinterpret_cast(chunk.data()); + while (remaining > 0) { + const unsigned int to_write = static_cast( + std::min(remaining, std::numeric_limits::max())); + const int written = gzwrite(file, data, to_write); + if (written == 0) { + int errnum = 0; + const char* error_message = gzerror(file, &errnum); + gzclose(file); + LOG(FATAL) << "Failed to write chunk: " << error_message; + } + data += written; + remaining -= static_cast(written); + } + + if (gzclose(file) != Z_OK) { + LOG(FATAL) << "Failed to close output file: " << output_path.string(); + } + + LOG(INFO) << "Wrote matching chunk to " << output_path.string(); +} + +void ProcessTar(const fs::path& tar_path, const fs::path& output_dir, + absl::Span expected_planes) { + std::unique_ptr source = std::make_unique( + tar_path, ChunkSourceLoaderConfig::V6TrainingData); + + const std::string base_name = tar_path.stem().string(); + size_t written_count = 0; + + for (size_t index = 0, total = source->GetChunkCount(); index < total; + ++index) { + const std::optional> chunk = + source->GetChunkData(index); + if (!chunk) { + LOG(WARNING) << "Skipping unreadable chunk " << index << " in " + << tar_path.string(); + continue; + } + + const std::optional match = + FindMatchingFrameIndex(*chunk, expected_planes); + if (!match) continue; + + WriteChunk(output_dir, base_name, index, *match, *chunk); + ++written_count; + } + + LOG(INFO) << "Finished processing " << tar_path.string() << ": wrote " + << written_count << " chunk(s)."; +} + +} // namespace + +int main(int argc, char** argv) { + absl::InitializeLog(); + absl::ParseCommandLine(argc, argv); + absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo); + + const fs::path input_dir(absl::GetFlag(FLAGS_input_dir)); + const fs::path output_dir(absl::GetFlag(FLAGS_output_dir)); + const std::string plane_values = absl::GetFlag(FLAGS_plane_values); + + if (!fs::exists(input_dir) || !fs::is_directory(input_dir)) { + LOG(FATAL) << "Input directory does not exist: " << input_dir.string(); + } + + const std::vector expected_planes = ParsePlaneValues(plane_values); + + fs::create_directories(output_dir); + + const std::vector tar_files = CollectTarFiles(input_dir); + const absl::Span expected_span(expected_planes); + + std::vector workers; + workers.reserve(tar_files.size()); + + for (const auto& tar_path : tar_files) { + workers.emplace_back([tar_path, output_dir, expected_span]() { + LOG(INFO) << "Processing tar file: " << tar_path.string(); + try { + ProcessTar(tar_path, output_dir, expected_span); + } catch (const std::exception& e) { + LOG(WARNING) << "Failed to process tar file " << tar_path << ": " + << e.what(); + } + }); + } + + for (auto& worker : workers) { + worker.join(); + } + + return 0; +} diff --git a/csrc/tools/position_weight_stats_main.cc b/csrc/tools/position_weight_stats_main.cc new file mode 100644 index 00000000..fd5dd2c4 --- /dev/null +++ b/csrc/tools/position_weight_stats_main.cc @@ -0,0 +1,213 @@ +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "loader/chunk_source/chunk_source.h" +#include "loader/chunk_source/tar_chunk_source.h" +#include "loader/stages/position_sampling.h" +#include "proto/data_loader_config.pb.h" +#include "trainingdata/trainingdata_v6.h" +#include "utils/training_data_printer.h" + +ABSL_FLAG(std::string, input_dir, "", "Directory to scan for .tar files."); +ABSL_FLAG(float, q_weight, 6.0, "Value for diff_focus_q_weight."); +ABSL_FLAG(float, pol_scale, 3.5, "Value for diff_focus_pol_scale."); + +namespace { + +namespace fs = std::filesystem; + +using ::lczero::training::ChunkSource; +using ::lczero::training::ChunkSourceLoaderConfig; +using ::lczero::training::ComputePositionSamplingWeight; +using ::lczero::training::FrameType; +using ::lczero::training::PositionSamplingConfig; +using ::lczero::training::PrintTrainingDataEntry; +using ::lczero::training::TarChunkSource; + +struct WeightedPosition { + FrameType data; + float weight; +}; + +std::vector CollectTarFiles(const fs::path& directory) { + std::vector files; + for (const auto& entry : fs::directory_iterator(directory)) { + const fs::path& path = entry.path(); + if (entry.is_regular_file() && path.extension() == ".tar") { + files.push_back(path); + } + } + absl::c_sort(files, [](const fs::path& lhs, const fs::path& rhs) { + return lhs.filename() < rhs.filename(); + }); + return files; +} + +std::vector CollectWeights(const fs::path& tar_path, + const PositionSamplingConfig& config, + WeightedPosition* max_weighted) { + std::vector weights; + std::unique_ptr source = std::make_unique( + tar_path, ChunkSourceLoaderConfig::V6TrainingData); + + const size_t total = source->GetChunkCount(); + for (size_t index = 0; index < total; ++index) { + if (index % 1000 == 0) { + LOG(INFO) << absl::StreamFormat(" Progress: %zu/%zu chunks (%.1f%%)", + index, total, 100.0 * index / total); + } + + const std::optional> chunk = + source->GetChunkData(index); + if (!chunk) { + LOG(WARNING) << "Skipping unreadable chunk " << index << " in " + << tar_path.string(); + continue; + } + + if (chunk->empty()) continue; + + for (const auto& entry : *chunk) { + const float weight = ComputePositionSamplingWeight(entry, config); + weights.push_back(weight); + if (max_weighted && weight > max_weighted->weight) { + max_weighted->data = entry; + max_weighted->weight = weight; + } + } + } + + return weights; +} + +void PrintHistogram(const std::vector& sorted_weights) { + if (sorted_weights.empty()) return; + + constexpr int kBuckets = 50; + constexpr int kMaxWidth = 60; + const float min_val = sorted_weights.front(); + const float max_val = sorted_weights.back(); + const float range = max_val - min_val; + + if (range == 0.0f) { + LOG(INFO) << "\nHistogram: All weights are identical (" << min_val << ")"; + return; + } + + std::vector buckets(kBuckets, 0); + for (float weight : sorted_weights) { + int bucket = static_cast((weight - min_val) / range * (kBuckets - 1)); + bucket = std::clamp(bucket, 0, kBuckets - 1); + ++buckets[bucket]; + } + + const int max_count = *absl::c_max_element(buckets); + if (max_count == 0) return; + + LOG(INFO) << "\nHistogram:"; + for (int bucket = 0; bucket < kBuckets; ++bucket) { + if (buckets[bucket] == 0) continue; + + const float bucket_start = min_val + range * bucket / kBuckets; + const float bucket_end = min_val + range * (bucket + 1) / kBuckets; + const int width = (buckets[bucket] * kMaxWidth + max_count / 2) / max_count; + + std::string bar; + for (int i = 0; i < width; ++i) bar += "█"; + + LOG(INFO) << absl::StreamFormat("[%.4f-%.4f) │%s (%d)", bucket_start, + bucket_end, bar, buckets[bucket]); + } +} + +void PrintPercentiles(const std::vector& sorted_weights) { + if (sorted_weights.empty()) return; + + LOG(INFO) << "\nPercentiles:"; + for (int p = 0; p <= 100; ++p) { + const size_t idx = (sorted_weights.size() - 1) * p / 100; + LOG(INFO) << absl::StreamFormat(" %3d%%: %.6f", p, sorted_weights[idx]); + } +} + +void PrintStatistics(const std::vector& weights) { + if (weights.empty()) { + LOG(INFO) << "No weights collected."; + return; + } + + const double sum = std::accumulate(weights.begin(), weights.end(), 0.0); + const double mean = sum / weights.size(); + + LOG(INFO) << "\nStatistics:"; + LOG(INFO) << " Total positions: " << weights.size(); + LOG(INFO) << absl::StreamFormat(" Mean: %.6f", mean); +} + +} // namespace + +int main(int argc, char** argv) { + absl::InitializeLog(); + absl::ParseCommandLine(argc, argv); + absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo); + + const std::string input_dir_str = absl::GetFlag(FLAGS_input_dir); + const float q_weight = absl::GetFlag(FLAGS_q_weight); + const float pol_scale = absl::GetFlag(FLAGS_pol_scale); + + if (input_dir_str.empty()) { + LOG(FATAL) << "--input_dir must be specified."; + } + + const fs::path input_dir(input_dir_str); + if (!fs::exists(input_dir) || !fs::is_directory(input_dir)) { + LOG(FATAL) << "Input directory does not exist: " << input_dir.string(); + } + + PositionSamplingConfig config; + config.set_diff_focus_q_weight(q_weight); + config.set_diff_focus_pol_scale(pol_scale); + + const std::vector tar_files = CollectTarFiles(input_dir); + LOG(INFO) << "Found " << tar_files.size() << " tar file(s)."; + + WeightedPosition max_weighted = {{}, 0.0f}; + std::vector all_weights; + for (const auto& tar_path : tar_files) { + LOG(INFO) << "Processing: " << tar_path.string(); + std::vector weights = + CollectWeights(tar_path, config, &max_weighted); + all_weights.insert(all_weights.end(), weights.begin(), weights.end()); + LOG(INFO) << " Collected " << weights.size() << " position(s)."; + } + + absl::c_sort(all_weights); + + PrintStatistics(all_weights); + PrintPercentiles(all_weights); + PrintHistogram(all_weights); + + if (max_weighted.weight > 0.0f) { + const std::string header = absl::StrFormat( + "\nPosition with highest weight (%.6f):", max_weighted.weight); + PrintTrainingDataEntry(max_weighted.data, header, 8, 4); + } + + return 0; +} diff --git a/csrc/tools/rescore_chunk_main.cc b/csrc/tools/rescore_chunk_main.cc new file mode 100644 index 00000000..a5ed374b --- /dev/null +++ b/csrc/tools/rescore_chunk_main.cc @@ -0,0 +1,166 @@ +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "chess/board.h" +#include "proto/data_loader_config.pb.h" +#include "syzygy/syzygy.h" +#include "trainingdata/reader.h" +#include "trainingdata/rescorer.h" +#include "trainingdata/trainingdata_v6.h" +#include "trainingdata/writer.h" +#include "utils/exception.h" + +ABSL_FLAG(std::string, chunk_path, "", + "Path to the chunk file (.gz) that should be rescored."); +ABSL_FLAG(std::string, syzygy_paths, "", + "Comma-separated list of Syzygy tablebase directories."); +ABSL_FLAG(double, dist_temp, 1.0, + "Policy temperature applied during rescoring."); +ABSL_FLAG(double, dist_offset, 0.0, "Policy offset applied during rescoring."); +ABSL_FLAG(double, dtz_boost, 0.0, + "DTZ boost applied during policy adjustments."); +ABSL_FLAG(int, new_input_format, -1, + "Optional conversion target for input format (-1 keeps original)."); +ABSL_FLAG( + double, deblunder_threshold, -1.0, + "Threshold for policy deblundering adjustments (negative to disable)."); +ABSL_FLAG( + double, deblunder_width, -1.0, + "Width controlling smoothing around threshold (negative to disable)."); + +namespace { + +namespace fs = std::filesystem; +using ::lczero::training::ChunkRescorerConfig; + +std::vector ReadChunkFrames(const fs::path& path) { + std::vector frames; + lczero::TrainingDataReader reader(path.string()); + lczero::V6TrainingData frame; + while (reader.ReadChunk(&frame)) { + frames.push_back(frame); + } + return frames; +} + +void WriteChunkFrames(const fs::path& path, + const std::vector& frames) { + lczero::TrainingDataWriter writer(path.string()); + for (const auto& frame : frames) { + writer.WriteChunk(frame); + } + writer.Finalize(); +} + +fs::path BuildOutputPath(const fs::path& input_path) { + fs::path directory = input_path.parent_path(); + fs::path stem = input_path.stem(); + fs::path filename = stem; + filename += "_rescored.gz"; + return directory / filename; +} + +} // namespace + +int main(int argc, char** argv) { + absl::InitializeLog(); + absl::ParseCommandLine(argc, argv); + absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo); + + const std::string chunk_path_flag = absl::GetFlag(FLAGS_chunk_path); + if (chunk_path_flag.empty()) { + LOG(FATAL) << "--chunk_path flag is required."; + } + const fs::path chunk_path(chunk_path_flag); + + ChunkRescorerConfig config; + + const std::string syzygy_paths_flag = absl::GetFlag(FLAGS_syzygy_paths); + if (!syzygy_paths_flag.empty()) { + config.set_syzygy_paths(syzygy_paths_flag); + } + config.set_dist_temp(static_cast(absl::GetFlag(FLAGS_dist_temp))); + config.set_dist_offset(static_cast(absl::GetFlag(FLAGS_dist_offset))); + config.set_dtz_boost(static_cast(absl::GetFlag(FLAGS_dtz_boost))); + config.set_new_input_format( + static_cast(absl::GetFlag(FLAGS_new_input_format))); + + const double deblunder_threshold_flag = + absl::GetFlag(FLAGS_deblunder_threshold); + const double deblunder_width_flag = absl::GetFlag(FLAGS_deblunder_width); + if (deblunder_threshold_flag >= 0.0 && deblunder_width_flag >= 0.0) { + config.set_deblunder_threshold( + static_cast(deblunder_threshold_flag)); + config.set_deblunder_width(static_cast(deblunder_width_flag)); + } else if (deblunder_threshold_flag >= 0.0 || deblunder_width_flag >= 0.0) { + LOG(FATAL) << "Both --deblunder_threshold and --deblunder_width must be " + << "set to non-negative values together."; + } + + LOG(INFO) << "Reading chunk from " << chunk_path.string(); + std::vector frames; + try { + frames = ReadChunkFrames(chunk_path); + } catch (const lczero::Exception& exception) { + LOG(FATAL) << "Failed to read chunk: " << exception.what(); + } + LOG(INFO) << "Loaded " << frames.size() << " frame(s) from chunk."; + if (frames.empty()) { + LOG(WARNING) << "Chunk contains no frames; writing empty output."; + try { + WriteChunkFrames(BuildOutputPath(chunk_path), frames); + } catch (const lczero::Exception& exception) { + LOG(FATAL) << "Failed to write rescored chunk: " << exception.what(); + } + return 0; + } + + lczero::InitializeMagicBitboards(); + + if (config.has_deblunder_threshold() && config.has_deblunder_width()) { + lczero::RescorerDeblunderSetup(config.deblunder_threshold(), + config.deblunder_width()); + } + + lczero::SyzygyTablebase tablebase; + if (!config.syzygy_paths().empty()) { + LOG(INFO) << "Initializing Syzygy tablebases from '" + << config.syzygy_paths() << "'."; + const std::string syzygy_paths(config.syzygy_paths()); + if (!tablebase.init(syzygy_paths)) { + LOG(WARNING) << "Failed to initialize Syzygy tablebases."; + } + } + + LOG(INFO) << "Rescoring chunk with dist_temp=" << config.dist_temp() + << ", dist_offset=" << config.dist_offset() + << ", dtz_boost=" << config.dtz_boost() + << ", new_input_format=" << config.new_input_format() << "."; + + try { + frames = lczero::RescoreTrainingData( + std::move(frames), &tablebase, config.dist_temp(), config.dist_offset(), + config.dtz_boost(), config.new_input_format()); + } catch (const lczero::Exception& exception) { + LOG(FATAL) << "Failed to rescore chunk: " << exception.what(); + } + + const fs::path output_path = BuildOutputPath(chunk_path); + LOG(INFO) << "Writing rescored chunk to " << output_path.string(); + try { + WriteChunkFrames(output_path, frames); + } catch (const lczero::Exception& exception) { + LOG(FATAL) << "Failed to write rescored chunk: " << exception.what(); + } + LOG(INFO) << "Completed rescoring of chunk."; + + return 0; +} diff --git a/csrc/tools/result_distribution_main.cc b/csrc/tools/result_distribution_main.cc new file mode 100644 index 00000000..5beaa718 --- /dev/null +++ b/csrc/tools/result_distribution_main.cc @@ -0,0 +1,208 @@ +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "loader/chunk_source/tar_chunk_source.h" +#include "trainingdata/trainingdata_v6.h" + +ABSL_FLAG(std::string, output_csv, "", + "Destination CSV file. Writes to stdout if empty."); +ABSL_FLAG(int, num_threads, 0, + "Number of worker threads. Defaults to hardware concurrency."); + +namespace { + +namespace fs = std::filesystem; + +using ::lczero::training::ChunkSourceLoaderConfig; +using ::lczero::training::FrameType; +using ::lczero::training::TarChunkSource; + +enum class ChunkResult { kWin, kDraw, kLoss }; + +struct ResultCounts { + uint64_t wins = 0; + uint64_t draws = 0; + uint64_t losses = 0; +}; + +class CsvWriter { + public: + CsvWriter(std::ostream* output, absl::Mutex* mutex) + : output_(output), mutex_(mutex) {} + + void Write(absl::string_view basename, const ResultCounts& counts) const { + absl::MutexLock lock(mutex_); + *output_ << basename << ',' << counts.wins << ',' << counts.draws << ',' + << counts.losses << '\n'; + output_->flush(); + } + + private: + std::ostream* output_; + absl::Mutex* mutex_; +}; + +std::ostream& SelectOutput(const fs::path& output_path, + std::ofstream& file_stream) { + if (output_path.empty()) return std::cout; + + file_stream.open(output_path, std::ios::out | std::ios::trunc); + if (!file_stream) { + LOG(FATAL) << "Failed to open output file: " << output_path.string(); + } + return file_stream; +} + +constexpr float kFloatTolerance = 1e-6f; + +std::optional DetermineChunkResult(absl::string_view chunk_payload, + size_t chunk_index, + const fs::path& tar_path) { + if (chunk_payload.size() < sizeof(FrameType)) { + LOG(WARNING) << "Chunk " << chunk_index << " in " << tar_path.string() + << " is too small."; + return std::nullopt; + } + + FrameType frame; + std::memcpy(&frame, chunk_payload.data(), sizeof(frame)); + + if (std::fabs(frame.result_d - 1.0f) <= kFloatTolerance) { + return ChunkResult::kDraw; + } + if (std::fabs(frame.result_d) > kFloatTolerance) { + LOG(WARNING) << "Chunk " << chunk_index << " in " << tar_path.string() + << " has unexpected result_d=" << frame.result_d << '.'; + return std::nullopt; + } + + const bool side_to_move = frame.side_to_move_or_enpassant != 0; + if (std::fabs(frame.result_q - 1.0f) <= kFloatTolerance) { + return side_to_move ? ChunkResult::kLoss : ChunkResult::kWin; + } + if (std::fabs(frame.result_q + 1.0f) <= kFloatTolerance) { + return side_to_move ? ChunkResult::kWin : ChunkResult::kLoss; + } + + LOG(WARNING) << "Chunk " << chunk_index << " in " << tar_path.string() + << " has unexpected result_q=" << frame.result_q << '.'; + return std::nullopt; +} + +ResultCounts CountResultsInTar(const fs::path& tar_path) { + ResultCounts counts; + + TarChunkSource source(tar_path, ChunkSourceLoaderConfig::V6TrainingData); + + const size_t chunk_count = source.GetChunkCount(); + for (size_t index = 0; index < chunk_count; ++index) { + const std::optional chunk = + source.GetChunkPrefix(index, sizeof(FrameType)); + if (!chunk) { + LOG(WARNING) << "Skipping unreadable chunk " << index << " in " + << tar_path.string(); + continue; + } + + const std::optional result = + DetermineChunkResult(*chunk, index, tar_path); + if (!result) continue; + + switch (*result) { + case ChunkResult::kWin: + ++counts.wins; + break; + case ChunkResult::kDraw: + ++counts.draws; + break; + case ChunkResult::kLoss: + ++counts.losses; + break; + } + } + + return counts; +} + +} // namespace + +int main(int argc, char** argv) { + absl::InitializeLog(); + std::vector positional = absl::ParseCommandLine(argc, argv); + absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo); + + if (positional.size() <= 1) { + LOG(FATAL) << "Provide at least one .tar file as a positional argument."; + } + + const fs::path output_path(absl::GetFlag(FLAGS_output_csv)); + std::ofstream file_stream; + std::ostream& output = SelectOutput(output_path, file_stream); + absl::Mutex output_mutex; + const CsvWriter writer(&output, &output_mutex); + + std::vector tar_files; + tar_files.reserve(positional.size() - 1); + for (size_t i = 1; i < positional.size(); ++i) { + tar_files.emplace_back(positional[i]); + } + + const int num_threads_flag = absl::GetFlag(FLAGS_num_threads); + size_t worker_count = 0; + if (num_threads_flag > 0) { + worker_count = static_cast(num_threads_flag); + } else { + const unsigned int hw_threads = std::thread::hardware_concurrency(); + worker_count = hw_threads > 0 ? static_cast(hw_threads) : 1; + } + worker_count = std::max(1, std::min(worker_count, tar_files.size())); + + std::atomic next_index(0); + std::vector workers; + workers.reserve(worker_count); + + for (size_t worker_id = 0; worker_id < worker_count; ++worker_id) { + workers.emplace_back([&tar_files, &next_index, &writer]() { + while (true) { + const size_t index = next_index.fetch_add(1, std::memory_order_relaxed); + if (index >= tar_files.size()) break; + const fs::path& tar_path = tar_files[index]; + LOG(INFO) << "Processing tar file: " << tar_path.string(); + try { + const ResultCounts counts = CountResultsInTar(tar_path); + writer.Write(tar_path.filename().string(), counts); + } catch (const std::exception& exception) { + LOG(WARNING) << "Failed to process tar file " << tar_path.string() + << ": " << exception.what(); + } + } + }); + } + + for (auto& worker : workers) { + worker.join(); + } + + return 0; +} diff --git a/csrc/tools/startpos_policy_distribution_main.cc b/csrc/tools/startpos_policy_distribution_main.cc new file mode 100644 index 00000000..a3de8c84 --- /dev/null +++ b/csrc/tools/startpos_policy_distribution_main.cc @@ -0,0 +1,146 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "loader/chunk_source/chunk_source.h" +#include "loader/chunk_source/tar_chunk_source.h" +#include "trainingdata/trainingdata_v6.h" + +ABSL_FLAG(std::string, input_dir, ".", "Directory to scan for .tar files."); +ABSL_FLAG(std::string, output_csv, "", + "Destination CSV file. Writes to stdout if empty."); + +namespace { + +namespace fs = std::filesystem; + +using ::lczero::training::ChunkSource; +using ::lczero::training::ChunkSourceLoaderConfig; +using ::lczero::training::FrameType; +using ::lczero::training::TarChunkSource; + +constexpr std::array kStartPositionPlanes = { + 0x000000000000ff00ull, 0x0000000000000042ull, 0x0000000000000024ull, + 0x0000000000000081ull, 0x0000000000000010ull, 0x0000000000000008ull, + 0x00ff000000000000ull, 0x4200000000000000ull, 0x2400000000000000ull, + 0x8100000000000000ull, 0x1000000000000000ull, 0x0800000000000000ull, + 0x0000000000000000ull, 0x0000000000000000ull, 0x0000000000000000ull, + 0x0000000000000000ull}; + +using PolicyProbe = std::pair; + +constexpr std::array kPolicyProbes = { + {{378, "g2g4"}, {346, "f2f3"}, {34, "b1a3"}, {161, "g1h3"}, + {403, "h2h4"}, {351, "f2f4"}, {234, "b2b4"}, {207, "a2a4"}, + {288, "d2d3"}, {204, "a2a3"}, {259, "c2c3"}, {36, "b1c3"}, + {400, "h2h3"}, {230, "b2b3"}, {322, "e2e4"}, {317, "e2e3"}, + {374, "g2g3"}, {264, "c2c4"}, {159, "g1f3"}, {293, "d2d4"}}}; + +bool MatchesStartPosition(const FrameType& data) { + return absl::c_equal( + kStartPositionPlanes, + absl::Span(data.planes, kStartPositionPlanes.size())); +} + +std::vector CollectTarFiles(const fs::path& directory) { + std::vector files; + for (const auto& entry : fs::directory_iterator(directory)) { + const fs::path& path = entry.path(); + if (entry.is_regular_file() && path.extension() == ".tar") { + files.push_back(path); + } + } + absl::c_sort(files, [](const fs::path& lhs, const fs::path& rhs) { + return lhs.filename() < rhs.filename(); + }); + return files; +} + +void WriteHeader(std::ostream& output) { + output << "file,index"; + for (const auto& probe : kPolicyProbes) output << ',' << probe.second; + output << '\n'; +} + +void WriteRow(std::ostream& output, absl::string_view sort_key, size_t index, + const FrameType& data) { + output << sort_key << ',' << index; + for (const auto& probe : kPolicyProbes) { + output << ',' << data.probabilities[probe.first]; + } + output << '\n'; +} + +void ProcessTarFile(const fs::path& tar_path, std::ostream& output) { + std::unique_ptr source = std::make_unique( + tar_path, ChunkSourceLoaderConfig::V6TrainingData); + const std::string sort_key = source->GetChunkSortKey(); + + for (size_t i = 0, count = source->GetChunkCount(); i < count; ++i) { + const std::optional> chunk = source->GetChunkData(i); + if (!chunk || chunk->empty()) continue; + + const FrameType& entry = chunk->front(); + if (!MatchesStartPosition(entry)) continue; + + WriteRow(output, sort_key, i, entry); + } +} + +std::ostream& SelectOutput(const fs::path& output_path, + std::ofstream& file_stream) { + if (output_path.empty()) return std::cout; + file_stream.open(output_path, std::ios::out | std::ios::trunc); + if (!file_stream) { + LOG(FATAL) << "Failed to open output file: " << output_path.string(); + } + return file_stream; +} + +} // namespace + +int main(int argc, char** argv) { + absl::InitializeLog(); + absl::ParseCommandLine(argc, argv); + + const fs::path input_dir(absl::GetFlag(FLAGS_input_dir)); + const fs::path output_path(absl::GetFlag(FLAGS_output_csv)); + + if (!fs::is_directory(input_dir)) { + LOG(FATAL) << "Input directory does not exist: " << input_dir.string(); + } + + std::ofstream file_stream; + std::ostream& output = SelectOutput(output_path, file_stream); + + WriteHeader(output); + + for (const auto& tar_path : CollectTarFiles(input_dir)) { + LOG(INFO) << "Processing tar file: " << tar_path.string(); + try { + ProcessTarFile(tar_path, output); + } catch (const std::exception& e) { + LOG(WARNING) << "Failed to process tar file " << tar_path << ": " + << e.what(); + } + } + + return 0; +} diff --git a/csrc/utils/gz.cc b/csrc/utils/gz.cc new file mode 100644 index 00000000..c03f4fc7 --- /dev/null +++ b/csrc/utils/gz.cc @@ -0,0 +1,51 @@ +#include "utils/gz.h" + +#include +#include + +#include +#include + +namespace lczero { +namespace training { + +std::string GunzipBuffer(std::string_view buffer) { + z_stream strm = {}; + int ret = inflateInit2(&strm, 16 + MAX_WBITS); + if (ret != Z_OK) { + throw GunzipError("Failed to initialize zlib inflate"); + } + + strm.avail_in = buffer.size(); + strm.next_in = reinterpret_cast(const_cast(buffer.data())); + + constexpr size_t kChunkSize = 16384; + std::string output; + std::array temp_buffer; + + do { + strm.avail_out = kChunkSize; + strm.next_out = reinterpret_cast(temp_buffer.data()); + + ret = inflate(&strm, Z_NO_FLUSH); + if (ret == Z_STREAM_ERROR || ret == Z_NEED_DICT || ret == Z_DATA_ERROR || + ret == Z_MEM_ERROR) { + inflateEnd(&strm); + throw GunzipError("zlib inflate error"); + } + + size_t bytes_written = kChunkSize - strm.avail_out; + output.append(temp_buffer.begin(), temp_buffer.begin() + bytes_written); + } while (strm.avail_out == 0); + + inflateEnd(&strm); + + if (ret != Z_STREAM_END) { + throw GunzipError("Incomplete gzip decompression"); + } + + return output; +} + +} // namespace training +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/gz.h b/csrc/utils/gz.h new file mode 100644 index 00000000..65ace23b --- /dev/null +++ b/csrc/utils/gz.h @@ -0,0 +1,18 @@ +#pragma once + +#include +#include +#include + +namespace lczero { +namespace training { + +class GunzipError : public std::runtime_error { + public: + using std::runtime_error::runtime_error; +}; + +std::string GunzipBuffer(std::string_view buffer); + +} // namespace training +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/metrics/exponential_aggregator.h b/csrc/utils/metrics/exponential_aggregator.h new file mode 100644 index 00000000..d7a73ee8 --- /dev/null +++ b/csrc/utils/metrics/exponential_aggregator.h @@ -0,0 +1,349 @@ +#pragma once + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace lczero { + +// 1 second period is exact. Other periods are powers of two. +enum class TimePeriod { + kEmpty = -128, + k1Millisecond = -10, + k2Milliseconds, + k4Milliseconds, + k8Milliseconds, + k16Milliseconds, + k31Milliseconds, + k63Milliseconds, + k125Milliseconds, + k250Milliseconds, + k500Milliseconds, + k1Second /* = 0 */, + k2Seconds, + k4Seconds, + k8Seconds, + k16Seconds, + k32Seconds, + k1Minute, + k2Minutes, + k4Minutes, + k9Minutes, + k17Minutes, + k36Minutes, + k1Hour, + k2Hours, + k5Hours, + k9Hours, + k18Hours, + k36Hours, + k3Days, + k6Days, + k12Days, + k24Days, + k49Days, + k97Days, + k194Days, + k388Days, + k777Days, + k2Years, + k4Years, + k9Years, + k17Years, /* = 29 */ + kAllTime = 127 +}; + +// ExponentialAggregator metrics over exponentially increasing time periods. +// +// The template parameter `Metric` can be used in two ways: +// 1. Function-based: Pass Clear and UpdateFrom functions to constructor. +// This is the new approach for protobuf-based metrics. +// 2. Method-based: The metric type has Reset() and MergeFrom() methods. +// This maintains backward compatibility with existing C++ struct metrics. +// +// For method-based metrics, the type must satisfy: +// * It must have a `Reset()` method that clears its state. +// * It must have a `MergeFrom(const Metric& other)` method to merge another +// metric into itself (used for bucket-to-bucket merging and live ingestion). +// * It must behave like a monoid (actually, unital magma is sufficient): +// * Merging with a default-constructed (empty) metric is a no-op. +// * The `MergeFrom` operation must be associative (actually, not really; +// currently we always merge old to new). It does not need to be +// commutative. +// * Having a `std::swap(Metric&, Metric&)` method is also beneficial. +template +class ExponentialAggregator { + public: + using Duration = std::chrono::nanoseconds; + using Clock = std::chrono::steady_clock; + using ClearFn = std::function; + using UpdateFromFn = std::function; + static constexpr TimePeriod kResolution = Resolution; + + // Resets the aggregator, clearing all buckets and pending metrics. + void Reset(Clock::time_point now = Clock::now()); + + // Ingests the passed metric into the pending bucket using MergeFrom + Reset. + void RecordMetrics(Metric&& metric); + + // Returns the latest completed metrics bucket for the given time period and + // duration since that period finished last time. If now is nullopt, it + // excludes the time since the last metrics flush. + std::pair GetBucketMetrics( + TimePeriod period, + std::optional now = std::nullopt) const; + + // Returns the current metrics that have been collected for at least the + // specified duration. Returns the metrics and the duration since the + // beginning of the covered period. If `now` is nullopt, it excludes metrics + // and the time since the last metrics flush. + std::pair GetAggregateEndingNow( + Duration duration, + std::optional now = Clock::now()) const; + + // Flushes the current pending bucket into the exponential metrics and + // advances time by the elapsed duration, potentially processing multiple + // ticks. Returns the largest time period that was updated by this advance + // (all smaller periods are also updated). + // + // Note: Advance() should typically be called more frequently than the tick + // frequency. When called less frequently (advancing multiple ticks at once), + // live statistics go into the first bucket and subsequent buckets are padded + // with empty metrics. + TimePeriod Advance(Clock::time_point now = Clock::now()); + + constexpr Duration GetResolution() const { return kPeriodDuration; } + + // Primary constructor for new protobuf-based metrics. + // Takes two free functions that define the metric's behavior. + ExponentialAggregator(ClearFn clear_fn, UpdateFromFn update_from_fn); + + // Default constructor for backward compatibility with old C++ metrics. + // This will only compile if `Metric` has Reset() and MergeFrom() methods. + ExponentialAggregator(); + + private: + // Constexpr power of 2 using bit shifts + static constexpr double constexpr_pow2(int exp) { + if (exp >= 0) { + return static_cast(1LL << exp); + } else { + return 1.0 / static_cast(1LL << (-exp)); + } + } + + static constexpr Duration kPeriodDuration = + std::chrono::duration_cast(std::chrono::duration( + constexpr_pow2(static_cast(Resolution)))); + + static size_t GetBucketIndex(TimePeriod period) { + return static_cast(period) - static_cast(Resolution); + } + + // The aggregation strategy is analogous to a binary counter. `tick_count_` + // represents the counter's value, and the `buckets_` array corresponds to its + // bits, each covering an exponentially larger time period. + // + // Advancing time increments `tick_count_`. When a bit flips from 1 to 0, + // its bucket's metric is merged (the "carry") into the next higher bucket. + // + // Note that buckets are never empty. A bucket whose corresponding bit in + // `tick_count_` is '0' simply holds the last complete metric for its time + // period. This ensures that a valid, historical metric is always available + // for the bucket query. + + ClearFn clear_fn_; + UpdateFromFn update_from_fn_; + + mutable absl::Mutex mutex_; + size_t tick_count_ ABSL_GUARDED_BY(mutex_); + // Buckets for each time period, starting from Resolution. + std::vector buckets_ ABSL_GUARDED_BY(mutex_); + Clock::time_point last_tick_time_ ABSL_GUARDED_BY(mutex_); + + mutable absl::Mutex pending_bucket_mutex_ ABSL_ACQUIRED_AFTER(mutex_); + Metric pending_bucket_ ABSL_GUARDED_BY(pending_bucket_mutex_); +}; + +template +ExponentialAggregator::ExponentialAggregator( + ClearFn clear_fn, UpdateFromFn update_from_fn) + : clear_fn_(std::move(clear_fn)), + update_from_fn_(std::move(update_from_fn)), + tick_count_(0), + last_tick_time_(Clock::now()) {} + +template +ExponentialAggregator::ExponentialAggregator() + : tick_count_(0), last_tick_time_(Clock::now()) { + // For backward compatibility, use metric methods if available + if constexpr (requires(Metric& m, const Metric& other) { + m.Reset(); + m.MergeFrom(other); + }) { + clear_fn_ = [](Metric& m) { m.Reset(); }; + update_from_fn_ = [](Metric& dest, const Metric& src) { + dest.MergeFrom(src); + }; + } else { + static_assert(sizeof(Metric) == 0, + "Metric type must have Reset() and MergeFrom() methods or " + "use function-based constructor"); + } +} + +template +void ExponentialAggregator::Reset( + std::chrono::steady_clock::time_point now) { + absl::MutexLock lock(&mutex_); + tick_count_ = 0; + buckets_.clear(); + last_tick_time_ = now; + + absl::MutexLock pending_bucket_lock(&pending_bucket_mutex_); + clear_fn_(pending_bucket_); +} + +template +void ExponentialAggregator::RecordMetrics(Metric&& metric) { + absl::MutexLock lock(&pending_bucket_mutex_); + update_from_fn_(pending_bucket_, metric); + clear_fn_(metric); +} + +template +auto ExponentialAggregator::GetBucketMetrics( + TimePeriod period, std::optional now) const + -> std::pair { + absl::MutexLock lock(&mutex_); + const size_t index = GetBucketIndex(period); + const Duration duration_since_update = + kPeriodDuration * (tick_count_ % (1ULL << index)) + + (now.has_value() ? Duration(*now - last_tick_time_) : Duration::zero()); + if (index >= buckets_.size()) return {Metric(), duration_since_update}; + return {buckets_[index], duration_since_update}; +} + +template +auto ExponentialAggregator::GetAggregateEndingNow( + Duration duration, std::optional now) const + -> std::pair { + Duration result_duration = Duration::zero(); + Metric result; + + { + absl::MutexLock lock(&mutex_); + if (now.has_value()) { + // If we'll use pending bucket, remove its duration from the request. + // The actual bucket we'll merge in the end as we have to merge newer + // into older buckets. + Duration duration_since_update = *now - last_tick_time_; + duration -= duration_since_update; + result_duration += duration_since_update; + } + + if (duration > Duration::zero()) { + // Convert the input `duration` into `num_ticks` (the number of base + // time periods), rounding up. + const auto div = std::div(duration.count(), kPeriodDuration.count()); + const size_t num_ticks = div.quot + bool(div.rem); + + // To cover the remaining `duration`, we select the minimal set of active + // historical buckets (where the corresponding bit in `tick_count_` is 1) + // that, when combined, meet or exceed the target duration. + // + // 1. First we determine the bit width of `num_ticks` (e.g., for 13 + // (1101b), the width is 4). Create a candidate set of ticks by + // masking `tick_count_` to this width. If this masked value is >= + // `num_ticks`, the corresponding set of buckets is sufficient. + // 2. If the masked value from step 2 is insufficient, we include one + // additional bucket. It doesn't matter if it's active or not. + size_t masked_ticks = [&]() { + const auto bit_width = std::bit_width(num_ticks); + const size_t mask = ~((~size_t{0}) << bit_width); + const size_t candidate_ticks = tick_count_ & mask; + if (candidate_ticks >= num_ticks) return candidate_ticks; + // One additional tick is needed. + return candidate_ticks | (size_t{1} << bit_width); + }(); + + while (masked_ticks) { + // Start merging from the highest bit (older bucket) to the lowest. + size_t idx = std::bit_width(masked_ticks) - 1; + masked_ticks &= ~(1ULL << idx); + if (idx < buckets_.size()) { + update_from_fn_(result, buckets_[idx]); + result_duration += kPeriodDuration * (1ULL << idx); + } + } + } + } + + if (now.has_value()) { + // The pending bucket is merged last, as we have to merge newer into older + // buckets. + absl::MutexLock lock(&pending_bucket_mutex_); + update_from_fn_(result, pending_bucket_); + } + + return {result, result_duration}; +} + +template +auto ExponentialAggregator::Advance(Clock::time_point now) + -> TimePeriod { + absl::MutexLock lock(&mutex_); + const int num_ticks_to_advance = (now - last_tick_time_) / kPeriodDuration; + if (num_ticks_to_advance <= 0) return TimePeriod::kEmpty; + + last_tick_time_ += num_ticks_to_advance * kPeriodDuration; + Metric live_carry; + { + // What was pending, now becomes carry. Pending bucket is cleared. + absl::MutexLock pending_bucket_lock(&pending_bucket_mutex_); + live_carry = std::move(pending_bucket_); + clear_fn_(pending_bucket_); + } + + const size_t initial_tick_count = tick_count_; + + auto one_tick = [&](Metric& carry) { + ++tick_count_; + + for (size_t i = 0;; ++i) { + const uint64_t interval_size = 1ULL << i; + if ((tick_count_ % interval_size) != 0) break; + while (i >= buckets_.size()) buckets_.emplace_back(); + // We always merge new into the old, so we swap the carry first, and then + // merge into it. + std::swap(carry, buckets_[i]); + update_from_fn_(carry, buckets_[i]); + } + }; + + // Carry the pending bucket into the first tick. + one_tick(live_carry); + // Then, if more than one tick is requested, we carry the empty bucket + // through the remaining ticks. + for (int i = 1; i < num_ticks_to_advance; ++i) { + Metric empty_carry; + one_tick(empty_carry); + } + + // Largest time period is the highest bit that was flipped in the process. + // To find it, we XOR the initial tick count with the current one, and + // find the highest bit. + return static_cast( + std::bit_width(initial_tick_count ^ tick_count_) - 1 + + static_cast(Resolution)); +} + +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/metrics/group.h b/csrc/utils/metrics/group.h new file mode 100644 index 00000000..d8c0c73e --- /dev/null +++ b/csrc/utils/metrics/group.h @@ -0,0 +1,99 @@ +#pragma once + +#include + +#include "utils/metrics/printer.h" + +namespace lczero { + +// Metric is a struct that implements the following interface: +// - void Reset(); // Resets the metric to its initial state. +// - void MergeFrom(const Metric& other); // Merges another metric into this +// one. Note that the incoming always happens later in time, so if e.g. merge +// keeps the latest value, it should update the current value with the incoming +// one. Used for bucket-to-bucket merging and live data ingestion. +// - (optional) std::string_view name() const; +// - (optional) std::string ToString() const; // If provided, returns a string +// representation of the metric. + +// Group several metric types together. This allows us to have a single +// `MetricGroup` that contains multiple different metrics. + +template +class MetricGroup { + public: + MetricGroup() = default; + + // Calls reset on all stats. + void Reset(); + + // Merges each individual stat from `other` into this group. + void MergeFrom(const MetricGroup& other); + + // Merges a single stat from `other` into this group. + template + void MergeFrom(const T& other); + + // Gets a const reference to a specific stat record. + template + const T& Get() const; + + // Gets a mutable pointer to a specific stat record. + template + T* GetMutable(); + + // Calls MetricPrinter for each stat in the group. + void Print(MetricPrinter& printer) const; + + private: + std::tuple stats_; +}; + +template +void MetricGroup::Reset() { + (std::get(stats_).Reset(), ...); +} + +template +void MetricGroup::MergeFrom( + const MetricGroup& other) { + (std::get(stats_).MergeFrom(std::get(other.stats_)), + ...); +} + +template +template +void MetricGroup::MergeFrom(const T& other) { + static_assert((std::is_same_v || ...), + "Type T must be one of the Stats types"); + std::get(stats_).MergeFrom(other); +} + +template +template +const T& MetricGroup::Get() const { + static_assert((std::is_same_v || ...), + "Type T must be one of the Stats types"); + return std::get(stats_); +} + +template +template +T* MetricGroup::GetMutable() { + static_assert((std::is_same_v || ...), + "Type T must be one of the Stats types"); + return &std::get(stats_); +} + +template +void MetricGroup::Print(MetricPrinter& printer) const { + ( + [&](const auto& stat) { + if constexpr (requires { stat.Print(printer); }) { + stat.Print(printer); + } + }(std::get(stats_)), + ...); +} + +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/metrics/load_metric.h b/csrc/utils/metrics/load_metric.h new file mode 100644 index 00000000..ca19520d --- /dev/null +++ b/csrc/utils/metrics/load_metric.h @@ -0,0 +1,114 @@ +#pragma once + +#include +#include + +#include +#include + +#include "proto/training_metrics.pb.h" + +namespace lczero { + +// ABOUTME: Helper class to manage timing logic for LoadMetricProto. +// ABOUTME: Tracks active periods and flushes accumulated time to the protobuf +// metric. +class LoadMetricUpdater { + public: + using Clock = std::chrono::steady_clock; + using Duration = std::chrono::duration; + + explicit LoadMetricUpdater(Clock::time_point initial_time = Clock::now()) + : last_flush_time_(initial_time), is_load_active_(true) {} + + // Starts tracking load from the given time point. + // Returns true if load was previously stopped (successful start). + bool LoadStart(Clock::time_point now = Clock::now()) { + absl::MutexLock lock(&mutex_); + FlushInternal(now); + bool was_stopped = !is_load_active_; + is_load_active_ = true; + return was_stopped; + } + + // Stops tracking load at the given time point. + // Returns true if load was previously active (successful stop). + bool LoadStop(Clock::time_point now = Clock::now()) { + absl::MutexLock lock(&mutex_); + FlushInternal(now); + bool was_active = is_load_active_; + is_load_active_ = false; + return was_active; + } + + // Flushes any uncounted load time into the metric. + void Flush(Clock::time_point now = Clock::now()) { + absl::MutexLock lock(&mutex_); + FlushInternal(now); + } + + // Flushes metrics and returns a copy, resetting the internal metric. + LoadMetricProto FlushMetrics(Clock::time_point now = Clock::now()) { + absl::MutexLock lock(&mutex_); + FlushInternal(now); + LoadMetricProto result = metric_; + metric_.Clear(); + return result; + } + + private: + // Flushes any uncounted load time into the metric (assumes mutex held). + void FlushInternal(Clock::time_point now) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + Duration elapsed = now - last_flush_time_; + double elapsed_seconds = elapsed.count(); + + metric_.set_total_seconds(metric_.total_seconds() + elapsed_seconds); + if (is_load_active_) { + metric_.set_load_seconds(metric_.load_seconds() + elapsed_seconds); + } + + last_flush_time_ = now; + } + + mutable absl::Mutex mutex_; + LoadMetricProto metric_ ABSL_GUARDED_BY(mutex_); + Clock::time_point last_flush_time_ ABSL_GUARDED_BY(mutex_); + bool is_load_active_ ABSL_GUARDED_BY(mutex_); +}; + +// UpdateFrom function for LoadMetricProto - simple additive behavior +inline void UpdateFrom(LoadMetricProto& dest, const LoadMetricProto& src) { + if (src.has_name()) dest.set_name(src.name()); + dest.set_load_seconds(dest.load_seconds() + src.load_seconds()); + dest.set_total_seconds(dest.total_seconds() + src.total_seconds()); +} + +// RAII class to temporarily pause load tracking. +class LoadMetricPauser { + public: + explicit LoadMetricPauser(LoadMetricUpdater& updater) : updater_(updater) { + successfully_paused_ = updater_.LoadStop(); + } + + ~LoadMetricPauser() { + if (successfully_paused_ && should_resume_) { + updater_.LoadStart(); + } + } + + // Prevents the pauser from resuming load tracking in destructor. + void DoNotResume() { should_resume_ = false; } + + LoadMetricPauser(const LoadMetricPauser&) = delete; + LoadMetricPauser& operator=(const LoadMetricPauser&) = delete; + LoadMetricPauser(LoadMetricPauser&&) = delete; + LoadMetricPauser& operator=(LoadMetricPauser&&) = delete; + + private: + LoadMetricUpdater& updater_; + bool successfully_paused_; + bool should_resume_ = true; +}; + +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/metrics/load_metric_test.cc b/csrc/utils/metrics/load_metric_test.cc new file mode 100644 index 00000000..ec1aebaa --- /dev/null +++ b/csrc/utils/metrics/load_metric_test.cc @@ -0,0 +1,325 @@ +#include "utils/metrics/load_metric.h" + +#include + +#include +#include + +#include "proto/training_metrics.pb.h" +#include "utils/metrics/exponential_aggregator.h" + +namespace lczero { +namespace training { + +class LoadMetricTest : public ::testing::Test { + protected: + using Clock = LoadMetricUpdater::Clock; + + void SetUp() override { start_time_ = Clock::now(); } + + Clock::time_point start_time_; +}; + +TEST_F(LoadMetricTest, BasicLoadMetricProto) { + LoadMetricProto metric; + EXPECT_EQ(metric.load_seconds(), 0.0); + EXPECT_EQ(metric.total_seconds(), 0.0); + + // Test UpdateFrom with LoadMetricUpdater + LoadMetricUpdater other_updater(start_time_); + other_updater.LoadStart(start_time_); + other_updater.LoadStop(start_time_ + std::chrono::milliseconds(500)); + LoadMetricProto other = + other_updater.FlushMetrics(start_time_ + std::chrono::milliseconds(500)); + UpdateFrom(metric, other); + EXPECT_NEAR(metric.load_seconds(), 0.5, 1e-6); + EXPECT_NEAR(metric.total_seconds(), 0.5, 1e-6); + + // Test Clear (used to be Reset) + metric.Clear(); + EXPECT_EQ(metric.load_seconds(), 0.0); + EXPECT_EQ(metric.total_seconds(), 0.0); +} + +TEST_F(LoadMetricTest, LoadMetricUpdaterBasic) { + LoadMetricUpdater updater(start_time_); + auto now = start_time_; + + // Start load and verify initial state + updater.LoadStart(now); + LoadMetricProto metric = updater.FlushMetrics(now); + EXPECT_EQ(metric.load_seconds(), 0.0); + EXPECT_EQ(metric.total_seconds(), 0.0); + + // Advance time and stop load + now += std::chrono::milliseconds(100); + updater.LoadStop(now); + metric = updater.FlushMetrics(now); + EXPECT_NEAR(metric.load_seconds(), 0.1, 1e-6); + EXPECT_NEAR(metric.total_seconds(), 0.1, 1e-6); + + // Wait idle time, then start again + now += std::chrono::milliseconds(50); + updater.LoadStart(now); + metric = updater.FlushMetrics(now); + EXPECT_NEAR(metric.load_seconds(), 0.0, 1e-6); // Reset after flush + EXPECT_NEAR(metric.total_seconds(), 0.05, 1e-6); // Only idle time + + now += std::chrono::milliseconds(200); + updater.LoadStop(now); + metric = updater.FlushMetrics(now); + EXPECT_NEAR(metric.load_seconds(), 0.2, 1e-6); // 0.2 load + EXPECT_NEAR(metric.total_seconds(), 0.2, 1e-6); // 0.2 total +} + +TEST_F(LoadMetricTest, LoadMetricUpdaterFlush) { + LoadMetricUpdater updater(start_time_); + auto now = start_time_; + + // Start load + updater.LoadStart(now); + now += std::chrono::milliseconds(100); + + // Flush should update the internal metric + updater.Flush(now); + LoadMetricProto metric = updater.FlushMetrics(now); + EXPECT_NEAR(metric.load_seconds(), 0.1, 1e-6); + EXPECT_NEAR(metric.total_seconds(), 0.1, 1e-6); + + // Continue loading + now += std::chrono::milliseconds(50); + updater.LoadStop(now); + metric = updater.FlushMetrics(now); + EXPECT_NEAR(metric.load_seconds(), 0.05, 1e-6); // Only new load time + EXPECT_NEAR(metric.total_seconds(), 0.05, 1e-6); // Only new total time +} + +TEST_F(LoadMetricTest, LoadMetricProtoMerging) { + LoadMetricUpdater updater1(start_time_); + LoadMetricUpdater updater2(start_time_); + auto now = start_time_; + + // Create load in updater1 + updater1.LoadStart(now); + updater1.LoadStop(now + std::chrono::milliseconds(100)); + LoadMetricProto metric1 = + updater1.FlushMetrics(now + std::chrono::milliseconds(100)); + EXPECT_NEAR(metric1.load_seconds(), 0.1, 1e-6); + EXPECT_NEAR(metric1.total_seconds(), 0.1, 1e-6); + + // Create load in updater2 + updater2.LoadStart(now); + updater2.LoadStop(now + std::chrono::milliseconds(100)); + LoadMetricProto metric2 = + updater2.FlushMetrics(now + std::chrono::milliseconds(100)); + EXPECT_NEAR(metric2.load_seconds(), 0.1, 1e-6); + EXPECT_NEAR(metric2.total_seconds(), 0.1, 1e-6); + + // Merge + UpdateFrom(metric1, metric2); + EXPECT_NEAR(metric1.load_seconds(), 0.2, 1e-6); + EXPECT_NEAR(metric1.total_seconds(), 0.2, 1e-6); + EXPECT_NEAR(metric2.load_seconds(), 0.1, 1e-6); // Source unchanged + EXPECT_NEAR(metric2.total_seconds(), 0.1, 1e-6); // Source unchanged +} + +TEST_F(LoadMetricTest, LoadMetricProtoMoveSemantics) { + // Test that LoadMetricProto move semantics work correctly + LoadMetricUpdater source_updater(start_time_); + source_updater.LoadStart(start_time_); + source_updater.LoadStop(start_time_ + std::chrono::milliseconds(100)); + LoadMetricProto source = + source_updater.FlushMetrics(start_time_ + std::chrono::milliseconds(100)); + EXPECT_NEAR(source.load_seconds(), 0.1, 1e-6); + EXPECT_NEAR(source.total_seconds(), 0.1, 1e-6); + + // Test move construction + LoadMetricProto moved_constructed(std::move(source)); + EXPECT_NEAR(moved_constructed.load_seconds(), 0.1, 1e-6); + EXPECT_NEAR(moved_constructed.total_seconds(), 0.1, 1e-6); + + // Test move assignment + LoadMetricProto move_assigned; + LoadMetricUpdater another_updater(start_time_); + another_updater.LoadStart(start_time_); + another_updater.LoadStop(start_time_ + std::chrono::milliseconds(50)); + LoadMetricProto another_source = + another_updater.FlushMetrics(start_time_ + std::chrono::milliseconds(50)); + EXPECT_NEAR(another_source.load_seconds(), 0.05, 1e-6); + EXPECT_NEAR(another_source.total_seconds(), 0.05, 1e-6); + + move_assigned = std::move(another_source); + EXPECT_NEAR(move_assigned.load_seconds(), 0.05, 1e-6); + EXPECT_NEAR(move_assigned.total_seconds(), 0.05, 1e-6); + + // Test UpdateFrom + LoadMetricProto dest; + UpdateFrom(dest, moved_constructed); + EXPECT_NEAR(dest.load_seconds(), 0.1, 1e-6); + EXPECT_NEAR(dest.total_seconds(), 0.1, 1e-6); + + UpdateFrom(dest, move_assigned); + EXPECT_NEAR(dest.load_seconds(), 0.15, 1e-6); + EXPECT_NEAR(dest.total_seconds(), 0.15, 1e-6); +} + +TEST_F(LoadMetricTest, LoadUtilizationTracking) { + LoadMetricUpdater updater(start_time_); + auto now = start_time_; + + // Start with some load time (load is active by default now) + now += std::chrono::milliseconds(100); + updater.LoadStop(now); // Stop load to create idle time + LoadMetricProto metric = updater.FlushMetrics(now); + EXPECT_NEAR(metric.load_seconds(), 0.1, 1e-6); // 100ms load + EXPECT_NEAR(metric.total_seconds(), 0.1, 1e-6); // 100ms total + + // Add some idle time (load is stopped) + now += std::chrono::milliseconds(100); + updater.LoadStart(now); // Start load again + metric = updater.FlushMetrics(now); + EXPECT_NEAR(metric.load_seconds(), 0.0, 1e-6); // No load time + EXPECT_NEAR(metric.total_seconds(), 0.1, 1e-6); // 100ms idle + + // Add more load time + now += std::chrono::milliseconds(200); + updater.LoadStop(now); + metric = updater.FlushMetrics(now); + EXPECT_NEAR(metric.load_seconds(), 0.2, 1e-6); // 200ms load + EXPECT_NEAR(metric.total_seconds(), 0.2, 1e-6); // 200ms total + + // Test complete utilization tracking with one updater + LoadMetricUpdater total_updater(start_time_); + auto total_now = start_time_; + + // 100ms load (active by default) + total_now += std::chrono::milliseconds(100); + total_updater.LoadStop(total_now); + + // 100ms idle + total_now += std::chrono::milliseconds(100); + total_updater.LoadStart(total_now); + + // 200ms load + total_now += std::chrono::milliseconds(200); + LoadMetricProto total_metric = total_updater.FlushMetrics(total_now); + + // Calculate utilization + double utilization = + total_metric.load_seconds() / total_metric.total_seconds(); + EXPECT_NEAR(utilization, 0.75, + 1e-6); // 75% utilization (300ms load / 400ms total) +} + +class LoadMetricProtoIntegrationTest : public ::testing::Test { + protected: + using TestAggregator = + ExponentialAggregator; + using Clock = TestAggregator::Clock; + + void SetUp() override { + aggregator_ = std::make_unique( + [](LoadMetricProto& m) { m.Clear(); }, + [](LoadMetricProto& dest, const LoadMetricProto& src) { + UpdateFrom(dest, src); + }); + start_time_ = Clock::now(); + } + + std::unique_ptr aggregator_; + Clock::time_point start_time_; +}; + +TEST_F(LoadMetricProtoIntegrationTest, RecordMetricsWithUpdater) { + auto current_time = start_time_; + + // Create metric with updater, simulate some load + LoadMetricUpdater updater(current_time); + updater.LoadStart(current_time); + current_time += std::chrono::milliseconds(150); + + // Flush and get metric + LoadMetricProto metric = updater.FlushMetrics(current_time); + + // Record the metric (this should use UpdateFrom + Reset) + aggregator_->RecordMetrics(std::move(metric)); + + // Get live metrics + auto [live_metrics, age] = aggregator_->GetAggregateEndingNow( + TestAggregator::Duration::zero(), current_time); + EXPECT_NEAR(live_metrics.load_seconds(), 0.15, 1e-6); +} + +TEST_F(LoadMetricProtoIntegrationTest, MultipleRecordMetrics) { + auto current_time = start_time_; + + // First metric + LoadMetricUpdater updater1(current_time); + updater1.LoadStart(current_time); + current_time += std::chrono::milliseconds(100); + LoadMetricProto metric1 = updater1.FlushMetrics(current_time); + aggregator_->RecordMetrics(std::move(metric1)); + + // Second metric + LoadMetricUpdater updater2(current_time); + updater2.LoadStart(current_time); + current_time += std::chrono::milliseconds(75); + LoadMetricProto metric2 = updater2.FlushMetrics(current_time); + aggregator_->RecordMetrics(std::move(metric2)); + + // Get live metrics + auto [live_metrics, age] = aggregator_->GetAggregateEndingNow( + TestAggregator::Duration::zero(), current_time); + EXPECT_NEAR(live_metrics.load_seconds(), 0.175, 1e-6); // 0.1 + 0.075 +} + +TEST_F(LoadMetricProtoIntegrationTest, AdvanceTest) { + auto current_time = start_time_; + + // Add some metrics + LoadMetricUpdater updater(current_time); + updater.LoadStart(current_time); + updater.LoadStop(current_time + std::chrono::milliseconds(100)); + LoadMetricProto metric = + updater.FlushMetrics(current_time + std::chrono::milliseconds(100)); + aggregator_->RecordMetrics(std::move(metric)); + + // Advance to move live metrics to buckets + auto tick_time = start_time_ + aggregator_->GetResolution(); + auto period = aggregator_->Advance(tick_time); + + // Should return the base time period + EXPECT_EQ(period, TimePeriod::k16Milliseconds); + + // Live metrics should be empty after advance + auto [live_metrics, age] = aggregator_->GetAggregateEndingNow( + TestAggregator::Duration::zero(), tick_time); + EXPECT_EQ(live_metrics.load_seconds(), 0.0); +} + +TEST_F(LoadMetricTest, LoadStartStopReturnValues) { + LoadMetricUpdater updater(start_time_); + + // Load starts active, so LoadStart should return false (already active) + EXPECT_FALSE(updater.LoadStart()); + + // LoadStop should return true (was active) + EXPECT_TRUE(updater.LoadStop()); + + // LoadStop again should return false (already stopped) + EXPECT_FALSE(updater.LoadStop()); + + // LoadStart should return true (was stopped) + EXPECT_TRUE(updater.LoadStart()); + + // LoadStart again should return false (already active) + EXPECT_FALSE(updater.LoadStart()); +} + +} // namespace training +} // namespace lczero + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} \ No newline at end of file diff --git a/csrc/utils/metrics/printer.h b/csrc/utils/metrics/printer.h new file mode 100644 index 00000000..2d56776c --- /dev/null +++ b/csrc/utils/metrics/printer.h @@ -0,0 +1,53 @@ +#pragma once + +#include + +#include +#include + +namespace lczero { + +class MetricPrinter { + public: + virtual ~MetricPrinter() = default; + virtual void StartGroup(std::string_view group_name) = 0; + virtual void Print(std::string_view metric_name, + const absl::AlphaNum& value) = 0; + virtual void EndGroup() = 0; +}; + +class StringMetricPrinter : public MetricPrinter { + public: + StringMetricPrinter(std::string* output) : output_(output) {} + void StartGroup(std::string_view group_name) override { + if (!first_group_) absl::StrAppend(output_, ", "); + absl::StrAppend(output_, group_name, "={"); + first_group_ = false; + first_metric_ = true; + } + + void Print(std::string_view metric_name, + const absl::AlphaNum& value) override { + if (!first_metric_) absl::StrAppend(output_, ", "); + absl::StrAppend(output_, metric_name, "=", value); + first_metric_ = false; + first_group_ = false; + } + + void EndGroup() override { absl::StrAppend(output_, "}"); } + + private: + std::string* output_; + bool first_metric_ = true; + bool first_group_ = true; +}; + +template +std::string MetricToString(const T& metric) { + std::string result; + StringMetricPrinter printer(&result); + metric.Print(printer); + return result; +} + +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/metrics/statistics_metric.h b/csrc/utils/metrics/statistics_metric.h new file mode 100644 index 00000000..de458b09 --- /dev/null +++ b/csrc/utils/metrics/statistics_metric.h @@ -0,0 +1,52 @@ +#pragma once + +#include + +#include "proto/training_metrics.pb.h" + +namespace lczero { + +// Helper function to add a sample to StatisticsProtoInt64 +inline void AddSample(StatisticsProtoInt64& stats, int64_t value) { + stats.set_min(std::min(stats.min(), value)); + stats.set_max(std::max(stats.max(), value)); + stats.set_sum(stats.sum() + value); + stats.set_count(stats.count() + 1); + stats.set_latest(value); +} + +// Helper function to add a sample to StatisticsProtoDouble +inline void AddSample(StatisticsProtoDouble& stats, double value) { + stats.set_min(std::min(stats.min(), value)); + stats.set_max(std::max(stats.max(), value)); + stats.set_sum(stats.sum() + value); + stats.set_count(stats.count() + 1); + stats.set_latest(value); +} + +// UpdateFrom function for StatisticsProtoInt64 - merges statistics +inline void UpdateFrom(StatisticsProtoInt64& dest, + const StatisticsProtoInt64& src) { + if (src.count() == 0) return; // Nothing to merge from empty source + + dest.set_min(std::min(dest.min(), src.min())); + dest.set_max(std::max(dest.max(), src.max())); + dest.set_sum(dest.sum() + src.sum()); + dest.set_count(dest.count() + src.count()); + dest.set_latest(src.latest()); // Source is newer, use its latest value +} + +// UpdateFrom function for StatisticsProtoDouble - merges statistics +inline void UpdateFrom(StatisticsProtoDouble& dest, + const StatisticsProtoDouble& src) { + if (src.count() == 0) return; // Nothing to merge from empty source + + if (src.has_name()) dest.set_name(src.name()); + dest.set_min(std::min(dest.min(), src.min())); + dest.set_max(std::max(dest.max(), src.max())); + dest.set_sum(dest.sum() + src.sum()); + dest.set_count(dest.count() + src.count()); + dest.set_latest(src.latest()); // Source is newer, use its latest value +} + +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/metrics/stats_test.cc b/csrc/utils/metrics/stats_test.cc new file mode 100644 index 00000000..9a58b747 --- /dev/null +++ b/csrc/utils/metrics/stats_test.cc @@ -0,0 +1,594 @@ +#include + +#include +#include +#include +#include + +#include "utils/metrics/exponential_aggregator.h" +#include "utils/metrics/group.h" +#include "utils/metrics/printer.h" + +namespace lczero { + +class CounterMetric { + public: + CounterMetric() : count_(0) {} + CounterMetric(int count) : count_(count) {} + + void Reset() { count_ = 0; } + + void MergeFrom(const CounterMetric& other) { count_ += other.count_; } + + void Print(MetricPrinter& printer) const { + printer.StartGroup("CounterMetric"); + printer.Print("count", static_cast(count_)); + printer.EndGroup(); + } + + int count() const { return count_; } + void set_count(int count) { count_ = count; } + + private: + int count_; +}; + +class AverageMetric { + public: + AverageMetric() : sum_(0), count_(0) {} + AverageMetric(double sum, int count) : sum_(sum), count_(count) {} + + void Reset() { + sum_ = 0; + count_ = 0; + } + + void MergeFrom(const AverageMetric& other) { + sum_ += other.sum_; + count_ += other.count_; + } + + void Print(MetricPrinter& printer) const { + printer.StartGroup("AverageMetric"); + printer.Print("sum", sum_); + printer.Print("count", static_cast(count_)); + if (count_ > 0) { + printer.Print("average", sum_ / count_); + } + printer.EndGroup(); + } + + double average() const { return count_ > 0 ? sum_ / count_ : 0.0; } + void add_sample(double value) { + sum_ += value; + count_++; + } + + double sum() const { return sum_; } + int count() const { return count_; } + + private: + double sum_; + int count_; +}; + +class MaxMetric { + public: + MaxMetric() : max_value_(0), has_value_(false) {} + MaxMetric(double max_value) : max_value_(max_value), has_value_(true) {} + + void Reset() { + max_value_ = 0; + has_value_ = false; + } + + void MergeFrom(const MaxMetric& other) { + if (other.has_value_) { + if (!has_value_ || other.max_value_ > max_value_) { + max_value_ = other.max_value_; + has_value_ = true; + } + } + } + + void Print(MetricPrinter& printer) const { + printer.StartGroup("MaxMetric"); + if (has_value_) { + printer.Print("max_value", max_value_); + printer.Print("has_value", static_cast(1)); + } else { + printer.Print("has_value", static_cast(0)); + } + printer.EndGroup(); + } + + double max_value() const { return max_value_; } + bool has_value() const { return has_value_; } + void set_value(double value) { + if (!has_value_ || value > max_value_) { + max_value_ = value; + has_value_ = true; + } + } + + private: + double max_value_; + bool has_value_; +}; + +// Optional value metric that demonstrates overshadowing behavior +class OptionalValueMetric { + public: + OptionalValueMetric() : value_(std::nullopt) {} + OptionalValueMetric(int value) : value_(value) {} + + void Reset() { value_ = std::nullopt; } + + void MergeFrom(const OptionalValueMetric& other) { + // Only copy the value if the other metric has one (overshadowing behavior) + if (other.value_.has_value()) { + value_ = other.value_; + } + } + + void Print(MetricPrinter& printer) const { + printer.StartGroup("OptionalValueMetric"); + if (value_.has_value()) { + printer.Print("value", value_.value()); + printer.Print("has_value", static_cast(1)); + } else { + printer.Print("has_value", static_cast(0)); + } + printer.EndGroup(); + } + + std::optional value() const { return value_; } + bool has_value() const { return value_.has_value(); } + void set_value(int value) { value_ = value; } + + private: + std::optional value_; +}; + +// Test MetricGroup functionality +class MetricGroupTest : public ::testing::Test { + protected: + using TestGroup = MetricGroup; + TestGroup group_; +}; + +TEST_F(MetricGroupTest, InitialState) { + // Test that metrics are initialized in their default state + EXPECT_EQ(group_.Get().count(), 0); + EXPECT_EQ(group_.Get().count(), 0); + EXPECT_FALSE(group_.Get().has_value()); +} + +TEST_F(MetricGroupTest, GetMutable) { + // Test getting mutable references and modifying them + auto* counter = group_.GetMutable(); + counter->set_count(42); + EXPECT_EQ(group_.Get().count(), 42); + + auto* average = group_.GetMutable(); + average->add_sample(10.0); + average->add_sample(20.0); + EXPECT_EQ(group_.Get().average(), 15.0); + + auto* max_metric = group_.GetMutable(); + max_metric->set_value(100.0); + EXPECT_EQ(group_.Get().max_value(), 100.0); +} + +TEST_F(MetricGroupTest, Reset) { + // Set up some data + group_.GetMutable()->set_count(42); + group_.GetMutable()->add_sample(10.0); + group_.GetMutable()->set_value(100.0); + + // Reset and verify everything is back to initial state + group_.Reset(); + + EXPECT_EQ(group_.Get().count(), 0); + EXPECT_EQ(group_.Get().count(), 0); + EXPECT_FALSE(group_.Get().has_value()); +} + +TEST_F(MetricGroupTest, MergeFromGroup) { + // Set up source group + TestGroup other; + other.GetMutable()->set_count(10); + other.GetMutable()->add_sample(5.0); + other.GetMutable()->set_value(50.0); + + // Set up destination group + group_.GetMutable()->set_count(20); + group_.GetMutable()->add_sample(15.0); + group_.GetMutable()->set_value(30.0); + + // Merge + group_.MergeFrom(other); + + // Verify results + EXPECT_EQ(group_.Get().count(), 30); // 20 + 10 + EXPECT_EQ(group_.Get().average(), 10.0); // (15 + 5) / 2 + EXPECT_EQ(group_.Get().max_value(), 50.0); // max(30, 50) +} + +TEST_F(MetricGroupTest, MergeFromSingleMetric) { + // Set up initial state + group_.GetMutable()->set_count(20); + + // Create a single metric to merge + CounterMetric counter(15); + + // Merge single metric + group_.MergeFrom(counter); + + // Verify result + EXPECT_EQ(group_.Get().count(), 35); // 20 + 15 +} + +TEST_F(MetricGroupTest, Print) { + // Set up data + group_.GetMutable()->set_count(42); + group_.GetMutable()->add_sample(10.0); + group_.GetMutable()->add_sample(20.0); + group_.GetMutable()->set_value(100.0); + + std::string result = MetricToString(group_); + + // Should contain all metric names and values + EXPECT_NE(result.find("CounterMetric"), std::string::npos); + EXPECT_NE(result.find("count=42"), std::string::npos); + EXPECT_NE(result.find("AverageMetric"), std::string::npos); + EXPECT_NE(result.find("average=15"), std::string::npos); // (10+20)/2 + EXPECT_NE(result.find("MaxMetric"), std::string::npos); + EXPECT_NE(result.find("max_value=100"), std::string::npos); +} + +// Test MetricToString functionality +class MetricPrinterTest : public ::testing::Test {}; + +TEST_F(MetricPrinterTest, StringMetricPrinter) { + std::string output; + StringMetricPrinter printer(&output); + + printer.StartGroup("test_group"); + printer.Print("metric1", std::string("value1")); + printer.Print("metric2", std::string("42")); + printer.EndGroup(); + + EXPECT_EQ(output, "test_group={metric1=value1, metric2=42}"); +} + +TEST_F(MetricPrinterTest, MultipleGroups) { + std::string output; + StringMetricPrinter printer(&output); + + printer.StartGroup("group1"); + printer.Print("metric1", std::string("value1")); + printer.EndGroup(); + + printer.StartGroup("group2"); + printer.Print("metric2", std::string("value2")); + printer.EndGroup(); + + EXPECT_EQ(output, "group1={metric1=value1}, group2={metric2=value2}"); +} + +TEST_F(MetricPrinterTest, EmptyGroup) { + std::string output; + StringMetricPrinter printer(&output); + + printer.StartGroup("empty_group"); + printer.EndGroup(); + + EXPECT_EQ(output, "empty_group={}"); +} + +TEST_F(MetricPrinterTest, SizeTOverload) { + std::string output; + StringMetricPrinter string_printer(&output); + MetricPrinter& printer = string_printer; // Use base class interface + + printer.StartGroup("test_group"); + printer.Print("count", static_cast(123)); + printer.EndGroup(); + + EXPECT_EQ(output, "test_group={count=123}"); +} + +TEST_F(MetricPrinterTest, MetricToStringFunction) { + CounterMetric counter(123); + std::string result = MetricToString(counter); + + EXPECT_NE(result.find("CounterMetric"), std::string::npos); + EXPECT_NE(result.find("123"), std::string::npos); +} + +class ExponentialAggregatorTest : public ::testing::Test { + protected: + using TestMetric = + MetricGroup; + using TestAggregator = + ExponentialAggregator; + + void SetUp() override { + // Create a fresh aggregator for each test to avoid state contamination + aggregator_ = std::make_unique(); + start_time_ = TestAggregator::Clock::now(); + aggregator_->Reset(start_time_); + } + + std::unique_ptr aggregator_; + TestAggregator::Clock::time_point start_time_; +}; + +TEST_F(ExponentialAggregatorTest, RecordMetrics) { + TestMetric metric; + metric.GetMutable()->set_count(10); + metric.GetMutable()->add_sample(5.0); + + // Update live metrics + aggregator_->RecordMetrics(std::move(metric)); + + // The original metric should be reset after move + EXPECT_EQ(metric.Get().count(), 0); + EXPECT_EQ(metric.Get().count(), 0); + + // Get live metrics to verify they were updated + auto [live_metrics, age] = + aggregator_->GetAggregateEndingNow(TestAggregator::Duration::zero()); + EXPECT_EQ(live_metrics.Get().count(), 10); + EXPECT_EQ(live_metrics.Get().average(), 5.0); +} + +TEST_F(ExponentialAggregatorTest, MultipleUpdatesLiveMetrics) { + // Update multiple times + for (int i = 1; i <= 5; ++i) { + TestMetric metric; + metric.GetMutable()->set_count(i); + metric.GetMutable()->add_sample(i * 2.0); + aggregator_->RecordMetrics(std::move(metric)); + } + + // Get live metrics + auto [live_metrics, age] = + aggregator_->GetAggregateEndingNow(TestAggregator::Duration::zero()); + EXPECT_EQ(live_metrics.Get().count(), 15); // 1+2+3+4+5 + EXPECT_EQ(live_metrics.Get().average(), + 6.0); // (2+4+6+8+10)/5 +} + +TEST_F(ExponentialAggregatorTest, Advance) { + // Add some live metrics + TestMetric metric; + metric.GetMutable()->set_count(10); + aggregator_->RecordMetrics(std::move(metric)); + + // Advance to move live metrics to buckets + auto tick_time = start_time_ + aggregator_->GetResolution(); + auto period = aggregator_->Advance(tick_time); + + // Should return the base time period + EXPECT_EQ(period, TimePeriod::k16Milliseconds); + + // Live metrics should be empty after tick + auto [live_metrics, age] = aggregator_->GetAggregateEndingNow( + TestAggregator::Duration::zero(), tick_time); + EXPECT_EQ(live_metrics.Get().count(), 0); +} + +TEST_F(ExponentialAggregatorTest, MultipleAdvances) { + // Add metrics and tick multiple times to test bucket management + auto current_time = start_time_; + for (const auto expected_period : { + TimePeriod::k16Milliseconds, + TimePeriod::k31Milliseconds, + TimePeriod::k16Milliseconds, + TimePeriod::k63Milliseconds, + TimePeriod::k16Milliseconds, + TimePeriod::k31Milliseconds, + TimePeriod::k16Milliseconds, + TimePeriod::k125Milliseconds, + }) { + TestMetric metric; + metric.GetMutable()->set_count(1); + aggregator_->RecordMetrics(std::move(metric)); + + current_time += aggregator_->GetResolution(); + auto period = aggregator_->Advance(current_time); + + EXPECT_EQ(period, expected_period); + } +} + +TEST_F(ExponentialAggregatorTest, MultipleAdvancesThreeTicks) { + // Add metrics and tick multiple times to test bucket management + auto current_time = start_time_; + for (const auto expected_period : { + TimePeriod::k31Milliseconds, + TimePeriod::k63Milliseconds, + TimePeriod::k125Milliseconds, + TimePeriod::k63Milliseconds, + }) { + TestMetric metric; + metric.GetMutable()->set_count(1); + aggregator_->RecordMetrics(std::move(metric)); + + current_time += aggregator_->GetResolution() * 3; + auto period = aggregator_->Advance(current_time); + + EXPECT_EQ(period, expected_period); + } +} + +TEST_F(ExponentialAggregatorTest, AggregationTest) { + TestMetric metric; + // We do 37 (0b100101) updates. + for (int i = 0; i < 37; ++i) { + metric.GetMutable()->set_count(i + 200); + metric.GetMutable()->set_value(i + 100); + aggregator_->RecordMetrics(std::move(metric)); + start_time_ += aggregator_->GetResolution(); + aggregator_->Advance(start_time_); + } + + // One more tick, but we don't advance this time. + metric.GetMutable()->set_count(1001); + metric.GetMutable()->set_value(1002); + aggregator_->RecordMetrics(std::move(metric)); + + const auto kRes = aggregator_->GetResolution(); + start_time_ += kRes / 3; // Not a tick yet. + + auto check_completed_bucket = [&](TimePeriod period, int expected_count, + std::optional expected_value, + TestAggregator::Duration expected_duration, + bool include_pending = false) { + auto [live_metrics, age] = aggregator_->GetBucketMetrics( + period, + include_pending ? std::make_optional(start_time_) : std::nullopt); + EXPECT_EQ(live_metrics.Get().count(), expected_count); + EXPECT_EQ(live_metrics.Get().has_value(), + expected_value.has_value()); + if (expected_value.has_value()) { + EXPECT_EQ(live_metrics.Get().value(), + expected_value.value()); + } + EXPECT_EQ(age, expected_duration); + }; + + check_completed_bucket(TimePeriod::k16Milliseconds, 236, 136, kRes * 0); + check_completed_bucket(TimePeriod::k31Milliseconds, 234 + 235, 135, kRes); + check_completed_bucket(TimePeriod::k63Milliseconds, 232 + 233 + 234 + 235, + 135, kRes); + check_completed_bucket(TimePeriod::k125Milliseconds, + 224 + 225 + 226 + 227 + 228 + 229 + 230 + 231, 131, + kRes * 5); + check_completed_bucket(TimePeriod::k125Milliseconds, + 224 + 225 + 226 + 227 + 228 + 229 + 230 + 231, 131, + kRes * 5 + kRes / 3, true); + check_completed_bucket(TimePeriod::k250Milliseconds, + 216 + 217 + 218 + 219 + 220 + 221 + 222 + 223 + 224 + + 225 + 226 + 227 + 228 + 229 + 230 + 231, + 131, kRes * 5); + check_completed_bucket(TimePeriod::k500Milliseconds, + 200 + 201 + 202 + 203 + 204 + 205 + 206 + 207 + 208 + + 209 + 210 + 211 + 212 + 213 + 214 + 215 + 216 + + 217 + 218 + 219 + 220 + 221 + 222 + 223 + 224 + + 225 + 226 + 227 + 228 + 229 + 230 + 231, + 131, kRes * 5); + check_completed_bucket(TimePeriod::k1Second, 0, std::nullopt, kRes * 37); + + auto check_aggregate = [&](TestAggregator::Duration duration, + int expected_count, + std::optional expected_value, + TestAggregator::Duration expected_duration, + bool include_pending = false) { + auto [live_metrics, age] = aggregator_->GetAggregateEndingNow( + duration, + include_pending ? std::make_optional(start_time_) : std::nullopt); + EXPECT_EQ(live_metrics.Get().count(), expected_count); + EXPECT_EQ(live_metrics.Get().has_value(), + expected_value.has_value()); + if (expected_value.has_value()) { + EXPECT_EQ(live_metrics.Get().value(), + expected_value.value()); + } + EXPECT_EQ(age, expected_duration); + }; + + check_aggregate(TestAggregator::Duration::zero(), 0, std::nullopt, kRes * 0); + check_aggregate(TestAggregator::Duration::zero(), 1001, 1002, kRes / 3, true); + check_aggregate(kRes / 4, 1001, 1002, kRes / 3, true); + check_aggregate(kRes / 10, 236, 136, kRes); + check_aggregate(kRes * 45 / 10, 232 + 233 + 234 + 235 + 236, 136, kRes * 5); + check_aggregate(kRes * 55 / 10, + 224 + 225 + 226 + 227 + 228 + 229 + 230 + 231 + 232 + 233 + + 234 + 235 + 236, + 136, kRes * (5 + 8)); +} + +TEST_F(ExponentialAggregatorTest, ActualVsRequestedTimeCoverage) { + // Test that GetAggregateEndingNow returns actual time covered by statistics + // rather than requested duration when insufficient historical data exists. + // This test recreates the scenario from the existing AggregationTest but + // specifically tests the requested vs actual duration behavior. + + const auto kRes = aggregator_->GetResolution(); + + // Set up aggregator with several data points like in AggregationTest + TestMetric metric; + for (int i = 0; i < 10; ++i) { + metric.GetMutable()->set_count(i + 200); + aggregator_->RecordMetrics(std::move(metric)); + start_time_ += kRes; + aggregator_->Advance(start_time_); + } + + // Based on AggregationTest line 507: request 4.5 * kRes, get back 5 * kRes + // This demonstrates that actual coverage (5 * kRes) can be MORE than + // requested (4.5 * kRes) because the aggregator only has specific bucket + // sizes available + const auto requested_duration = kRes * 45 / 10; // 4.5 * kRes + auto [result_metrics, actual_duration] = + aggregator_->GetAggregateEndingNow(requested_duration, std::nullopt); + + // The key test: when requesting 4.5 * kRes, we should get actual time covered + // which may be different than the requested amount due to bucket granularity + EXPECT_GT(actual_duration, requested_duration); // Actual > requested + EXPECT_GT(actual_duration, std::chrono::nanoseconds::zero()); + + // Verify we got some metrics (non-zero count) + EXPECT_GT(result_metrics.Get().count(), 0); + + // Test that shows the key behavior: when we request more time than available, + // we get back only the time that's actually covered by data + auto [result_zero, duration_zero] = aggregator_->GetAggregateEndingNow( + kRes * 100, std::nullopt); // Request way more + + // The returned duration should be much less than requested (showing actual vs + // requested) + const auto huge_request = kRes * 100; + EXPECT_LT(duration_zero, huge_request); + EXPECT_GT(duration_zero, std::chrono::nanoseconds::zero()); +} + +TEST_F(ExponentialAggregatorTest, ExactDurationTest) { + // Simple test: add buckets for exactly 5 seconds, request kAllTime, + // ensure we get back exactly 5.0 seconds duration (not more) + + auto current_time = start_time_; + + // Add buckets for exactly 5 seconds + for (int i = 0; i < 5; ++i) { + TestMetric metric; + metric.GetMutable()->set_count(100 + i); + aggregator_->RecordMetrics(std::move(metric)); + current_time += std::chrono::seconds(1); + aggregator_->Advance(current_time); + } + + // Request statistics for all time + auto [result_metrics, actual_duration] = aggregator_->GetAggregateEndingNow( + std::chrono::duration_cast( + std::chrono::hours(24 * 365)), // Request way more than 5 seconds + std::nullopt); + + // Should return exactly 5.0 seconds duration (actual time covered) + const auto expected_duration = std::chrono::seconds(5); + EXPECT_EQ(actual_duration, expected_duration); + + // Should have all our data + const int expected_total = 100 + 101 + 102 + 103 + 104; + EXPECT_EQ(result_metrics.Get().count(), expected_total); +} + +} // namespace lczero + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} \ No newline at end of file diff --git a/csrc/utils/queue.h b/csrc/utils/queue.h new file mode 100644 index 00000000..bfd558f9 --- /dev/null +++ b/csrc/utils/queue.h @@ -0,0 +1,697 @@ +#pragma once + +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/fixed_array.h" +#include "absl/log/log.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" + +namespace lczero { + +// Virtual base class for type-erased handling of queues. +class QueueBase { + public: + virtual ~QueueBase() = default; + virtual size_t Size() const = 0; + virtual size_t Capacity() const = 0; + virtual bool IsClosed() const = 0; + virtual void Close() = 0; + virtual size_t GetTotalPutCount(bool reset = false) = 0; + virtual size_t GetTotalGetCount(bool reset = false) = 0; + virtual size_t GetTotalDropCount(bool reset = false) = 0; +}; + +// Exception thrown when queue operations are attempted on a closed queue. +class QueueClosedException : public std::runtime_error { + public: + QueueClosedException() : std::runtime_error("Queue is closed") {} +}; + +// Exception thrown when queue operation is cancelled via stop_token. +class QueueRequestCancelled : public std::runtime_error { + public: + QueueRequestCancelled() : std::runtime_error("Queue request cancelled") {} +}; + +enum class OverflowBehavior { BLOCK, DROP_NEW, KEEP_NEWEST }; + +// Thread-safe fixed-size circular buffer queue with blocking operations. +// Supports both single and batch put/get operations. +// The queue automatically closes when all Producer tokens are destroyed. +// When closed, Put operations throw immediately, but Get operations only throw +// when the queue becomes empty - allowing consumption of remaining elements. +template +class Queue : public QueueBase { + public: + // Backwards-compatible alias to support code referring to + // Queue::OverflowBehavior. + using OverflowBehavior = ::lczero::OverflowBehavior; + + explicit Queue(size_t capacity, + OverflowBehavior overflow_behavior = OverflowBehavior::BLOCK); + + // RAII token for producers. Queue automatically closes when all producers + // are destroyed. All Put operations must go through this class. + class Producer { + public: + explicit Producer(Queue& queue); + ~Producer(); + + // Move constructor and assignment + Producer(Producer&& other) noexcept; + Producer& operator=(Producer&& other) noexcept; + + // Disable copy to maintain RAII semantics + Producer(const Producer&) = delete; + Producer& operator=(const Producer&) = delete; + + // Puts a single element into the queue. Blocks if queue is full. + void Put(const T& item, std::stop_token stop_token = {}); + void Put(T&& item, std::stop_token stop_token = {}); + + // Puts multiple elements into the queue. Blocks if not enough space. + void Put(absl::Span items, std::stop_token stop_token = {}); + void Put(absl::Span items, std::stop_token stop_token = {}); + + // Explicitly close this producer, decrementing the producer count + void Close(); + + private: + Queue* queue_; + }; + + // Creates a new producer token for this queue. + Producer CreateProducer(); + + // Gets a single element from the queue. Blocks if queue is empty. + T Get(std::stop_token stop_token = {}); + + // Gets exactly count elements from the queue. Blocks until count elements + // available. + absl::FixedArray Get(size_t count, std::stop_token stop_token = {}); + + // Gets a single element from the queue if available, returns std::nullopt + // if empty. + std::optional MaybeGet(); + + // Returns the current size of the queue. + size_t Size() const override; + + // Returns the capacity of the queue. + size_t Capacity() const override; + + // Explicitly close the queue, preventing further Put operations. + void Close() override; + + // Returns true if the queue is closed. + bool IsClosed() const override; + + // Wait until queue has at least the specified amount of free space. + void WaitForRoomAtLeast(size_t room, std::stop_token stop_token = {}); + + // Wait until queue has at most the specified amount of free space. + void WaitForRoomAtMost(size_t room, std::stop_token stop_token = {}); + + // Wait until queue has at least the specified number of elements. + void WaitForSizeAtLeast(size_t size, std::stop_token stop_token = {}); + + // Wait until queue has at most the specified number of elements. + void WaitForSizeAtMost(size_t size, std::stop_token stop_token = {}); + + // Returns the total number of elements that have been put into the queue. + // If reset is true, resets the counter to 0 after returning the value. + size_t GetTotalPutCount(bool reset = false) override; + + // Returns the total number of elements that have been retrieved from the + // queue. If reset is true, resets the counter to 0 after returning the value. + size_t GetTotalGetCount(bool reset = false) override; + + // Returns the total number of elements that have been dropped from the queue. + // If reset is true, resets the counter to 0 after returning the value. + size_t GetTotalDropCount(bool reset = false) override; + + private: + friend class Producer; + + const size_t capacity_; + const OverflowBehavior overflow_behavior_; + absl::FixedArray buffer_ ABSL_GUARDED_BY(mutex_); + size_t head_ ABSL_GUARDED_BY(mutex_) = 0; + size_t tail_ ABSL_GUARDED_BY(mutex_) = 0; + size_t size_ ABSL_GUARDED_BY(mutex_) = 0; + size_t producer_count_ ABSL_GUARDED_BY(mutex_) = 0; + bool closed_ ABSL_GUARDED_BY(mutex_) = false; + size_t total_put_count_ ABSL_GUARDED_BY(mutex_) = 0; + size_t total_get_count_ ABSL_GUARDED_BY(mutex_) = 0; + size_t total_drop_count_ ABSL_GUARDED_BY(mutex_) = 0; + + mutable absl::Mutex mutex_; + absl::CondVar cond_var_; + + // Internal methods for producer management + void RemoveProducer(); + + // Internal Put methods (called by Producer) + void PutInternal(const T& item, std::stop_token stop_token = {}); + void PutInternal(T&& item, std::stop_token stop_token = {}); + void PutInternal(absl::Span items, std::stop_token stop_token = {}); + void PutInternal(absl::Span items, std::stop_token stop_token = {}); + + // Condition predicates for blocking operations + bool CanPutOne() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + bool CanGet() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + // Additional condition predicates for wait functions + bool HasRoomAtLeast(size_t room) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + bool HasRoomAtMost(size_t room) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + bool HasSizeAtLeast(size_t size) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + bool HasSizeAtMost(size_t size) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); +}; + +// Implementation + +template +Queue::Queue(size_t capacity, OverflowBehavior overflow_behavior) + : capacity_(capacity), + overflow_behavior_(overflow_behavior), + buffer_(capacity) {} + +// Producer implementation +template +Queue::Producer::Producer(Queue& queue) : queue_(&queue) { + // Producer count is incremented in CreateProducer() + VLOG(1) << "Queue@" << static_cast(queue_) << " producer@" + << static_cast(this) << " constructed."; +} + +template +Queue::Producer::~Producer() { + if (queue_) { + VLOG(1) << "Queue@" << static_cast(queue_) << " producer@" + << static_cast(this) << " destructing."; + queue_->RemoveProducer(); + } +} + +template +Queue::Producer::Producer(Producer&& other) noexcept : queue_(other.queue_) { + other.queue_ = nullptr; +} + +template +typename Queue::Producer& Queue::Producer::operator=( + Producer&& other) noexcept { + if (this != &other) { + if (queue_) { + queue_->RemoveProducer(); + } + queue_ = other.queue_; + other.queue_ = nullptr; + } + return *this; +} + +template +void Queue::Producer::Put(const T& item, std::stop_token stop_token) { + queue_->PutInternal(item, stop_token); +} + +template +void Queue::Producer::Put(T&& item, std::stop_token stop_token) { + queue_->PutInternal(std::move(item), stop_token); +} + +template +void Queue::Producer::Put(absl::Span items, + std::stop_token stop_token) { + queue_->PutInternal(items, stop_token); +} + +template +void Queue::Producer::Put(absl::Span items, std::stop_token stop_token) { + queue_->PutInternal(items, stop_token); +} + +template +void Queue::Producer::Close() { + if (queue_) { + VLOG(1) << "Queue@" << static_cast(queue_) << " producer@" + << static_cast(this) << " close invoked."; + queue_->RemoveProducer(); + queue_ = nullptr; + } +} + +// Queue implementation +template +typename Queue::Producer Queue::CreateProducer() { + absl::MutexLock lock(&mutex_); + if (closed_) throw QueueClosedException(); + ++producer_count_; + return Producer(*this); +} + +template +void Queue::RemoveProducer() { + absl::MutexLock lock(&mutex_); + --producer_count_; + if (producer_count_ == 0 && !closed_) { + closed_ = true; + VLOG(1) << "Queue@" << static_cast(this) + << " closed after last producer removed."; + cond_var_.SignalAll(); + } +} + +template +void Queue::PutInternal(const T& item, std::stop_token stop_token) { + absl::MutexLock lock(&mutex_); + if (closed_) { + VLOG(1) << "Queue@" << static_cast(this) + << " PutInternal(const&) throwing QueueClosedException;" + << " producers=" << producer_count_; + throw QueueClosedException(); + } + ++total_put_count_; + + switch (overflow_behavior_) { + case OverflowBehavior::BLOCK: { + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!CanPutOne()) { + if (closed_) throw QueueClosedException(); + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } + if (closed_) throw QueueClosedException(); + break; + } + case OverflowBehavior::DROP_NEW: + if (size_ >= capacity_) { + ++total_drop_count_; + return; + } + break; + case OverflowBehavior::KEEP_NEWEST: + if (size_ >= capacity_) { + head_ = (head_ + 1) % capacity_; + --size_; + ++total_drop_count_; + } + break; + } + + buffer_[tail_] = item; + tail_ = (tail_ + 1) % capacity_; + ++size_; + cond_var_.SignalAll(); +} + +template +void Queue::PutInternal(T&& item, std::stop_token stop_token) { + absl::MutexLock lock(&mutex_); + if (closed_) { + VLOG(1) << "Queue@" << static_cast(this) + << " PutInternal(T&&) throwing QueueClosedException;" + << " producers=" << producer_count_; + throw QueueClosedException(); + } + ++total_put_count_; + + switch (overflow_behavior_) { + case OverflowBehavior::BLOCK: { + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!CanPutOne()) { + if (closed_) throw QueueClosedException(); + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } + if (closed_) throw QueueClosedException(); + break; + } + case OverflowBehavior::DROP_NEW: + if (size_ >= capacity_) { + ++total_drop_count_; + return; + } + break; + case OverflowBehavior::KEEP_NEWEST: + if (size_ >= capacity_) { + head_ = (head_ + 1) % capacity_; + --size_; + ++total_drop_count_; + } + break; + } + + buffer_[tail_] = std::move(item); + tail_ = (tail_ + 1) % capacity_; + ++size_; + cond_var_.SignalAll(); +} + +template +void Queue::PutInternal(absl::Span items, + std::stop_token stop_token) { + if (items.empty()) return; + + size_t remaining = items.size(); + size_t offset = 0; + + while (remaining > 0) { + absl::MutexLock lock(&mutex_); + if (closed_) { + VLOG(1) << "Queue@" << static_cast(this) + << " PutInternal(span const) throwing QueueClosedException;" + << " producers=" << producer_count_; + throw QueueClosedException(); + } + + size_t batch_size; + switch (overflow_behavior_) { + case OverflowBehavior::BLOCK: { + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!CanPutOne()) { + if (closed_) throw QueueClosedException(); + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } + if (closed_) throw QueueClosedException(); + batch_size = std::min(remaining, capacity_ - size_); + break; + } + case OverflowBehavior::DROP_NEW: + batch_size = std::min(remaining, capacity_ - size_); + if (batch_size == 0) { + total_put_count_ += remaining; + total_drop_count_ += remaining; + return; + } + break; + case OverflowBehavior::KEEP_NEWEST: + batch_size = std::min(remaining, capacity_); + while (size_ + batch_size > capacity_) { + head_ = (head_ + 1) % capacity_; + --size_; + ++total_drop_count_; + } + break; + } + + for (size_t i = 0; i < batch_size; ++i) { + buffer_[tail_] = items[offset + i]; + tail_ = (tail_ + 1) % capacity_; + ++size_; + } + total_put_count_ += batch_size; + cond_var_.SignalAll(); + + offset += batch_size; + remaining -= batch_size; + } +} + +template +void Queue::PutInternal(absl::Span items, std::stop_token stop_token) { + if (items.empty()) return; + + size_t remaining = items.size(); + size_t offset = 0; + + while (remaining > 0) { + absl::MutexLock lock(&mutex_); + if (closed_) { + VLOG(1) << "Queue@" << static_cast(this) + << " PutInternal(span) throwing QueueClosedException;" + << " producers=" << producer_count_; + throw QueueClosedException(); + } + + size_t batch_size; + switch (overflow_behavior_) { + case OverflowBehavior::BLOCK: { + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!CanPutOne()) { + if (closed_) throw QueueClosedException(); + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } + if (closed_) throw QueueClosedException(); + batch_size = std::min(remaining, capacity_ - size_); + break; + } + case OverflowBehavior::DROP_NEW: + batch_size = std::min(remaining, capacity_ - size_); + if (batch_size == 0) { + total_put_count_ += remaining; + total_drop_count_ += remaining; + return; + } + break; + case OverflowBehavior::KEEP_NEWEST: + batch_size = std::min(remaining, capacity_); + while (size_ + batch_size > capacity_) { + head_ = (head_ + 1) % capacity_; + --size_; + ++total_drop_count_; + } + break; + } + + for (size_t i = 0; i < batch_size; ++i) { + buffer_[tail_] = std::move(items[offset + i]); + tail_ = (tail_ + 1) % capacity_; + ++size_; + } + total_put_count_ += batch_size; + cond_var_.SignalAll(); + + offset += batch_size; + remaining -= batch_size; + } +} + +template +T Queue::Get(std::stop_token stop_token) { + absl::MutexLock lock(&mutex_); + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!CanGet()) { + if (closed_ && size_ == 0) { + VLOG(1) << "Queue@" << static_cast(this) + << " Get() throwing QueueClosedException; producers=" + << producer_count_; + throw QueueClosedException(); + } + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } + if (closed_ && size_ == 0) { + VLOG(1) << "Queue@" << static_cast(this) + << " Get() throwing QueueClosedException; producers=" + << producer_count_; + throw QueueClosedException(); + } + + T item = std::move(buffer_[head_]); + head_ = (head_ + 1) % capacity_; + --size_; + ++total_get_count_; + cond_var_.SignalAll(); + + return item; +} + +template +absl::FixedArray Queue::Get(size_t count, std::stop_token stop_token) { + if (count == 0) return absl::FixedArray(0); + + absl::FixedArray result(count); + size_t remaining = count; + size_t offset = 0; + + while (remaining > 0) { + absl::MutexLock lock(&mutex_); + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!CanGet()) { + if (closed_ && size_ == 0) { + VLOG(1) << "Queue@" << static_cast(this) << " Get(" + << count << ") throwing QueueClosedException; producers=" + << producer_count_; + throw QueueClosedException(); + } + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } + if (closed_ && size_ == 0) { + VLOG(1) << "Queue@" << static_cast(this) << " Get(" << count + << ") throwing QueueClosedException; producers=" + << producer_count_; + throw QueueClosedException(); + } + + size_t batch_size = std::min(remaining, size_); + + for (size_t i = 0; i < batch_size; ++i) { + result[offset + i] = std::move(buffer_[head_]); + head_ = (head_ + 1) % capacity_; + --size_; + ++total_get_count_; + } + cond_var_.SignalAll(); + + offset += batch_size; + remaining -= batch_size; + } + + return result; +} + +template +std::optional Queue::MaybeGet() { + absl::MutexLock lock(&mutex_); + if (size_ == 0) return std::nullopt; + + T item = std::move(buffer_[head_]); + head_ = (head_ + 1) % capacity_; + --size_; + ++total_get_count_; + cond_var_.SignalAll(); + + return item; +} + +template +size_t Queue::Size() const { + absl::MutexLock lock(&mutex_); + return size_; +} + +template +size_t Queue::Capacity() const { + return capacity_; +} + +template +void Queue::Close() { + absl::MutexLock lock(&mutex_); + if (!closed_) { + closed_ = true; + VLOG(1) << "Queue@" << static_cast(this) + << " closed explicitly; producers=" << producer_count_; + cond_var_.SignalAll(); + } +} + +template +bool Queue::IsClosed() const { + absl::MutexLock lock(&mutex_); + return closed_; +} + +template +void Queue::WaitForRoomAtLeast(size_t room, std::stop_token stop_token) { + absl::MutexLock lock(&mutex_); + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!HasRoomAtLeast(room)) { + if (closed_) throw QueueClosedException(); + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } +} + +template +void Queue::WaitForRoomAtMost(size_t room, std::stop_token stop_token) { + absl::MutexLock lock(&mutex_); + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!HasRoomAtMost(room)) { + if (closed_) throw QueueClosedException(); + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } +} + +template +void Queue::WaitForSizeAtLeast(size_t size, std::stop_token stop_token) { + absl::MutexLock lock(&mutex_); + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!HasSizeAtLeast(size)) { + if (closed_) throw QueueClosedException(); + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } +} + +template +void Queue::WaitForSizeAtMost(size_t size, std::stop_token stop_token) { + absl::MutexLock lock(&mutex_); + std::stop_callback cb(stop_token, [this]() { cond_var_.SignalAll(); }); + while (!HasSizeAtMost(size)) { + if (closed_) throw QueueClosedException(); + if (stop_token.stop_requested()) throw QueueRequestCancelled(); + cond_var_.Wait(&mutex_); + } +} + +template +bool Queue::CanPutOne() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + return closed_ || size_ < capacity_; +} + +template +bool Queue::CanGet() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + return closed_ || size_ > 0; +} + +template +bool Queue::HasRoomAtLeast(size_t room) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + return capacity_ - size_ >= room; +} + +template +bool Queue::HasRoomAtMost(size_t room) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + return capacity_ - size_ <= room; +} + +template +bool Queue::HasSizeAtLeast(size_t size) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + return size_ >= size; +} + +template +bool Queue::HasSizeAtMost(size_t size) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + return size_ <= size; +} + +template +size_t Queue::GetTotalPutCount(bool reset) { + absl::MutexLock lock(&mutex_); + size_t count = total_put_count_; + if (reset) total_put_count_ = 0; + return count; +} + +template +size_t Queue::GetTotalGetCount(bool reset) { + absl::MutexLock lock(&mutex_); + size_t count = total_get_count_; + if (reset) total_get_count_ = 0; + return count; +} + +template +size_t Queue::GetTotalDropCount(bool reset) { + absl::MutexLock lock(&mutex_); + size_t count = total_drop_count_; + if (reset) total_drop_count_ = 0; + return count; +} + +} // namespace lczero diff --git a/csrc/utils/queue_test.cc b/csrc/utils/queue_test.cc new file mode 100644 index 00000000..19520edd --- /dev/null +++ b/csrc/utils/queue_test.cc @@ -0,0 +1,1602 @@ +// ABOUTME: Comprehensive unit tests for the Queue template class +// ABOUTME: Tests thread-safe operations, blocking behavior, and edge cases + +#include "utils/queue.h" + +#include + +#include +#include +#include +#include +#include + +namespace lczero { + +class QueueTest : public ::testing::Test { + protected: + void SetUp() override {} +}; + +// Basic functionality tests + +TEST_F(QueueTest, ConstructorCreatesEmptyQueue) { + Queue queue(5); + EXPECT_EQ(queue.Size(), 0); + EXPECT_EQ(queue.Capacity(), 5); +} + +TEST_F(QueueTest, SinglePutGet) { + Queue queue(5); + { + auto producer = queue.CreateProducer(); + producer.Put(42); + EXPECT_EQ(queue.Size(), 1); + } // Producer destroyed here, queue closes + + int value = queue.Get(); + EXPECT_EQ(value, 42); + EXPECT_EQ(queue.Size(), 0); +} + +TEST_F(QueueTest, MovePutGet) { + Queue> queue(5); + { + auto producer = queue.CreateProducer(); + auto ptr = std::make_unique(42); + producer.Put(std::move(ptr)); + EXPECT_EQ(queue.Size(), 1); + } // Producer destroyed here, queue closes + + auto result = queue.Get(); + EXPECT_EQ(*result, 42); + EXPECT_EQ(queue.Size(), 0); +} + +TEST_F(QueueTest, MultiplePutGet) { + Queue queue(5); + + { + auto producer = queue.CreateProducer(); + for (int i = 0; i < 5; ++i) { + producer.Put(i); + } + EXPECT_EQ(queue.Size(), 5); + } // Producer destroyed here, queue closes + + for (int i = 0; i < 5; ++i) { + int value = queue.Get(); + EXPECT_EQ(value, i); + } + EXPECT_EQ(queue.Size(), 0); +} + +TEST_F(QueueTest, CircularBufferBehavior) { + Queue queue(3); + auto producer = queue.CreateProducer(); + + // Fill queue + producer.Put(1); + producer.Put(2); + producer.Put(3); + + // Get one item, put another + EXPECT_EQ(queue.Get(), 1); + producer.Put(4); + + // Verify remaining items + EXPECT_EQ(queue.Get(), 2); + EXPECT_EQ(queue.Get(), 3); + EXPECT_EQ(queue.Get(), 4); +} + +// Batch operations tests + +TEST_F(QueueTest, BatchPutConstSpan) { + Queue queue(5); + std::vector items = {1, 2, 3}; + + { + auto producer = queue.CreateProducer(); + producer.Put(absl::Span(items)); + EXPECT_EQ(queue.Size(), 3); + } // Producer destroyed here, queue closes + + for (int i = 0; i < 3; ++i) { + EXPECT_EQ(queue.Get(), i + 1); + } +} + +TEST_F(QueueTest, BatchPutMoveSpan) { + Queue> queue(5); + std::vector> items; + items.push_back(std::make_unique(1)); + items.push_back(std::make_unique(2)); + items.push_back(std::make_unique(3)); + + { + auto producer = queue.CreateProducer(); + producer.Put(absl::Span>(items)); + EXPECT_EQ(queue.Size(), 3); + } // Producer destroyed here, queue closes + + for (int i = 0; i < 3; ++i) { + auto result = queue.Get(); + EXPECT_EQ(*result, i + 1); + } +} + +TEST_F(QueueTest, BatchPutEmptySpan) { + Queue queue(5); + std::vector empty_items; + + { + auto producer = queue.CreateProducer(); + producer.Put(absl::Span(empty_items)); + EXPECT_EQ(queue.Size(), 0); + } // Producer destroyed here, queue closes +} + +TEST_F(QueueTest, BatchGet) { + Queue queue(5); + + { + auto producer = queue.CreateProducer(); + for (int i = 0; i < 5; ++i) { + producer.Put(i); + } + } // Producer destroyed here, queue closes + + auto result = queue.Get(3); + EXPECT_EQ(result.size(), 3); + for (int i = 0; i < 3; ++i) { + EXPECT_EQ(result[i], i); + } + EXPECT_EQ(queue.Size(), 2); +} + +TEST_F(QueueTest, BatchGetZeroCount) { + Queue queue(5); + { + auto producer = queue.CreateProducer(); + producer.Put(42); + } // Producer destroyed here, queue closes + + auto result = queue.Get(0); + EXPECT_EQ(result.size(), 0); + EXPECT_EQ(queue.Size(), 1); +} + +// Edge cases and error conditions + +TEST_F(QueueTest, CapacityOne) { + Queue queue(1); + + { + auto producer = queue.CreateProducer(); + producer.Put(42); + EXPECT_EQ(queue.Size(), 1); + } // Producer destroyed here, queue closes + + EXPECT_EQ(queue.Get(), 42); + EXPECT_EQ(queue.Size(), 0); +} + +// Tests for operations when all producer tokens are destroyed + +TEST_F(QueueTest, CreateProducerOnClosedQueue) { + Queue queue(5); + // Create and immediately destroy producer to close queue + { + auto producer = queue.CreateProducer(); + } + + // Trying to create a new producer after queue is closed results in an + // exception. + EXPECT_THROW(queue.CreateProducer(), QueueClosedException); +} + +TEST_F(QueueTest, GetOnClosedQueue) { + Queue queue(5); + // Create and immediately destroy producer to close queue + { + auto producer = queue.CreateProducer(); + } + + EXPECT_THROW(queue.Get(), QueueClosedException); +} + +TEST_F(QueueTest, BatchGetOnClosedQueue) { + Queue queue(5); + // Create and immediately destroy producer to close queue + { + auto producer = queue.CreateProducer(); + } + + EXPECT_THROW(queue.Get(3), QueueClosedException); +} + +// Thread safety tests + +TEST_F(QueueTest, SingleProducerSingleConsumer) { + Queue queue(10); + std::atomic producer_done{false}; + std::vector consumed; + + std::thread producer([&queue, &producer_done]() { + auto prod = queue.CreateProducer(); + for (int i = 0; i < 100; ++i) { + prod.Put(i); + } + producer_done = true; + // Producer destroyed here, closing the queue + }); + + std::thread consumer([&queue, &consumed, &producer_done]() { + int value; + while (!producer_done || queue.Size() > 0) { + try { + value = queue.Get(); + consumed.push_back(value); + } catch (const QueueClosedException&) { + break; + } + } + }); + + producer.join(); + consumer.join(); + + EXPECT_EQ(consumed.size(), 100); + for (int i = 0; i < 100; ++i) { + EXPECT_EQ(consumed[i], i); + } +} + +TEST_F(QueueTest, MultipleProducersMultipleConsumers) { + Queue queue(10); + constexpr int num_producers = 2; + constexpr int items_per_producer = 5; + constexpr int total_items = num_producers * items_per_producer; + + std::vector all_consumed; + std::vector producers; + + // Use a single producer token that we control explicitly + auto producer_token = queue.CreateProducer(); + + // Start producers - they all share the same producer token via reference + for (int p = 0; p < num_producers; ++p) { + producers.emplace_back([&producer_token, p]() { + for (int i = 0; i < items_per_producer; ++i) { + int value = p * items_per_producer + i; + producer_token.Put(value); + } + }); + } + + // Wait for all producers to finish + for (auto& producer : producers) { + producer.join(); + } + + // Now explicitly close the queue by destroying the producer token + { + auto temp = std::move(producer_token); + } // Queue is now closed + + // Now consume all items from the closed queue + for (int i = 0; i < total_items; ++i) { + all_consumed.push_back(queue.Get()); + } + + // Verify all items were consumed + EXPECT_EQ(all_consumed.size(), total_items); + EXPECT_EQ(queue.Size(), 0); + + // Trying to get one more should throw + EXPECT_THROW(queue.Get(), QueueClosedException); +} + +TEST_F(QueueTest, BlockingBehaviorOnFullQueue) { + Queue queue(2); + std::promise about_to_block; + std::future about_to_block_future = about_to_block.get_future(); + std::atomic put_completed{false}; + auto producer = queue.CreateProducer(); + + // Fill the queue + producer.Put(1); + producer.Put(2); + + std::thread blocker([&producer, &about_to_block, &put_completed]() { + about_to_block.set_value(); // Signal we're about to block + producer.Put(3); // This should block + put_completed = true; + }); + + // Wait for thread to signal it's about to block + about_to_block_future.wait(); + EXPECT_FALSE(put_completed); + + // Make space in the queue + EXPECT_EQ(queue.Get(), 1); + + blocker.join(); + EXPECT_TRUE(put_completed); + EXPECT_EQ(queue.Size(), 2); +} + +TEST_F(QueueTest, BlockingBehaviorOnEmptyQueue) { + Queue queue(5); + std::promise about_to_block; + std::future about_to_block_future = about_to_block.get_future(); + std::atomic get_completed{false}; + std::atomic result{-1}; + auto producer = queue.CreateProducer(); + + std::thread blocker([&queue, &about_to_block, &get_completed, &result]() { + about_to_block.set_value(); // Signal we're about to block + result = queue.Get(); // This should block + get_completed = true; + }); + + // Wait for thread to signal it's about to block + about_to_block_future.wait(); + EXPECT_FALSE(get_completed); + + // Put an item in the queue + producer.Put(42); + + blocker.join(); + EXPECT_TRUE(get_completed); + EXPECT_EQ(result, 42); +} + +TEST_F(QueueTest, ProducerDestructionUnblocksWaitingGet) { + Queue queue(5); // Empty queue + + std::promise about_to_block; + std::future about_to_block_future = about_to_block.get_future(); + std::atomic exception_thrown{false}; + + // Create a producer to keep queue open initially + std::unique_ptr::Producer> producer = + std::make_unique::Producer>(queue.CreateProducer()); + + std::thread blocker([&queue, &about_to_block, &exception_thrown]() { + about_to_block.set_value(); // Signal we're about to block + try { + queue.Get(); // This should block + } catch (const QueueClosedException&) { + exception_thrown = true; + } + }); + + // Wait for thread to signal it's about to block + about_to_block_future.wait(); + EXPECT_FALSE(exception_thrown); + + // Destroy the producer - this should close queue and unblock the waiting + // Get() + producer.reset(); + + blocker.join(); + EXPECT_TRUE(exception_thrown); +} + +// Test: Get() should not throw when queue is closed but has elements +TEST_F(QueueTest, GetFromClosedQueueWithElements) { + Queue queue(5); + + // Put some elements in the queue, then destroy producer to close it + { + auto producer = queue.CreateProducer(); + producer.Put(1); + producer.Put(2); + producer.Put(3); + EXPECT_EQ(queue.Size(), 3); + } // Producer destroyed here, queue closes + + // Should be able to get elements that were already in the queue + EXPECT_EQ(queue.Get(), 1); + EXPECT_EQ(queue.Get(), 2); + EXPECT_EQ(queue.Get(), 3); + EXPECT_EQ(queue.Size(), 0); + + // Only now should Get() throw when queue is empty and closed + EXPECT_THROW(queue.Get(), QueueClosedException); +} + +TEST_F(QueueTest, BatchGetFromClosedQueueWithElements) { + Queue queue(5); + + // Put some elements in the queue, then destroy producer to close it + { + auto producer = queue.CreateProducer(); + producer.Put(1); + producer.Put(2); + producer.Put(3); + EXPECT_EQ(queue.Size(), 3); + } // Producer destroyed here, queue closes + + // Should be able to get elements that were already in the queue + auto result = queue.Get(2); + EXPECT_EQ(result.size(), 2); + EXPECT_EQ(result[0], 1); + EXPECT_EQ(result[1], 2); + EXPECT_EQ(queue.Size(), 1); + + // Get remaining element + EXPECT_EQ(queue.Get(), 3); + EXPECT_EQ(queue.Size(), 0); + + // Only now should Get() throw when queue is empty and closed + EXPECT_THROW(queue.Get(1), QueueClosedException); +} + +// Test producer token mechanism specifically +TEST_F(QueueTest, ProducerTokenMechanism) { + Queue queue(5); + + // Create multiple producers + auto producer1 = queue.CreateProducer(); + auto producer2 = queue.CreateProducer(); + + // Both should be able to put items + producer1.Put(1); + producer2.Put(2); + EXPECT_EQ(queue.Size(), 2); + + // Destroy one producer - queue should still be open + { + auto temp = std::move(producer1); + } // producer1 is destroyed here + producer2.Put(3); + EXPECT_EQ(queue.Size(), 3); + + // Destroy last producer - queue should close + { + auto temp = std::move(producer2); + } // producer2 is destroyed here + + // Should still be able to get existing items + EXPECT_EQ(queue.Get(), 1); + EXPECT_EQ(queue.Get(), 2); + EXPECT_EQ(queue.Get(), 3); + + // But trying to get more should throw + EXPECT_THROW(queue.Get(), QueueClosedException); +} + +TEST_F(QueueTest, ProducerMoveSemantics) { + Queue queue(5); + + auto producer1 = queue.CreateProducer(); + producer1.Put(42); + + // Move constructor + auto producer2 = std::move(producer1); + producer2.Put(43); + EXPECT_EQ(queue.Size(), 2); + + // Create another producer and use move assignment + auto producer3 = queue.CreateProducer(); + producer3 = std::move(producer2); + producer3.Put(44); + EXPECT_EQ(queue.Size(), 3); + + // Destroy the last producer + { + auto temp = std::move(producer3); + } // producer3 is destroyed here + + // Should be able to get all items + EXPECT_EQ(queue.Get(), 42); + EXPECT_EQ(queue.Get(), 43); + EXPECT_EQ(queue.Get(), 44); + EXPECT_THROW(queue.Get(), QueueClosedException); +} + +// Tests for Put operations on closed queue +TEST_F(QueueTest, PutOnClosedQueueThrowsException) { + Queue queue(5); + + // Create producer and close it + auto producer = queue.CreateProducer(); + queue.Close(); + + // All Put operations should throw on closed queue + EXPECT_THROW(producer.Put(42), QueueClosedException); + EXPECT_THROW(producer.Put(std::move(42)), QueueClosedException); + + std::vector items = {1, 2, 3}; + EXPECT_THROW(producer.Put(absl::Span(items)), + QueueClosedException); + EXPECT_THROW(producer.Put(absl::Span(items)), QueueClosedException); +} + +TEST_F(QueueTest, PutOnClosedQueueAfterProducerDestruction) { + Queue queue(5); + + // Create producer, add item, then close by destroying all producers + auto producer = queue.CreateProducer(); + producer.Put(1); + { + auto temp_producer = std::move(producer); + } // All producers destroyed, queue closed + + // Try to create new producer after close + EXPECT_THROW(queue.CreateProducer(), QueueClosedException); +} + +TEST_F(QueueTest, BatchPutOnClosedQueueThrowsException) { + Queue queue(10); + + auto producer = queue.CreateProducer(); + queue.Close(); + + // Batch put operations should throw on closed queue + std::vector items = {1, 2, 3, 4, 5}; + EXPECT_THROW(producer.Put(absl::Span(items)), + QueueClosedException); + + std::vector mutable_items = {6, 7, 8}; + EXPECT_THROW(producer.Put(absl::Span(mutable_items)), + QueueClosedException); +} + +TEST_F(QueueTest, PublicCloseMethod) { + Queue queue(5); + + auto producer = queue.CreateProducer(); + producer.Put(1); + producer.Put(2); + + // Explicitly close the queue using public Close() method + queue.Close(); + + // Put operations should now throw + EXPECT_THROW(producer.Put(3), QueueClosedException); + + // But Get operations should still work for existing items + EXPECT_EQ(queue.Get(), 1); + EXPECT_EQ(queue.Get(), 2); + + // Get should throw when queue is empty and closed + EXPECT_THROW(queue.Get(), QueueClosedException); +} + +TEST_F(QueueTest, CloseUnblocksWaitingSinglePut) { + Queue queue(2); // Small capacity + auto producer = queue.CreateProducer(); + + // Fill the queue + producer.Put(1); + producer.Put(2); + + std::promise about_to_block; + std::future about_to_block_future = about_to_block.get_future(); + std::atomic exception_thrown{false}; + + std::thread blocker([&producer, &about_to_block, &exception_thrown]() { + about_to_block.set_value(); // Signal we're about to block + try { + producer.Put(3); // This should block since queue is full + } catch (const QueueClosedException&) { + exception_thrown = true; + } + }); + + // Wait for thread to signal it's about to block + about_to_block_future.wait(); + EXPECT_FALSE(exception_thrown); + + // Close the queue - this should unblock the waiting Put() + queue.Close(); + + blocker.join(); + EXPECT_TRUE(exception_thrown); +} + +TEST_F(QueueTest, CloseUnblocksWaitingBatchPut) { + Queue queue(3); // Small capacity + auto producer = queue.CreateProducer(); + + // Fill the queue partially + producer.Put(1); + producer.Put(2); + + std::promise about_to_block; + std::future about_to_block_future = about_to_block.get_future(); + std::atomic exception_thrown{false}; + + std::thread blocker([&producer, &about_to_block, &exception_thrown]() { + about_to_block.set_value(); // Signal we're about to block + try { + std::vector items = {3, 4, 5}; // Need 3 slots but only 1 available + producer.Put(absl::Span(items)); // This should block + } catch (const QueueClosedException&) { + exception_thrown = true; + } + }); + + // Wait for thread to signal it's about to block + about_to_block_future.wait(); + EXPECT_FALSE(exception_thrown); + + // Close the queue - this should unblock the waiting batch Put() + queue.Close(); + + blocker.join(); + EXPECT_TRUE(exception_thrown); +} + +// Tests for new wait functions +TEST_F(QueueTest, WaitForRoomAtLeast) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + // Initially empty queue should have room >= 5 + queue.WaitForRoomAtLeast(5); + EXPECT_EQ(queue.Size(), 0); + + // Fill queue partially + producer.Put(1); + producer.Put(2); + + // Should have room >= 3 + queue.WaitForRoomAtLeast(3); + EXPECT_EQ(queue.Size(), 2); + + // Fill more + producer.Put(3); + producer.Put(4); + + // Should have room >= 1 + queue.WaitForRoomAtLeast(1); + EXPECT_EQ(queue.Size(), 4); + + // Test blocking behavior + producer.Put(5); // Queue is now full + + std::promise wait_started; + std::future wait_started_future = wait_started.get_future(); + std::atomic wait_completed{false}; + + std::thread waiter([&queue, &wait_started, &wait_completed]() { + wait_started.set_value(); + queue.WaitForRoomAtLeast(2); // Should block until 2 slots are free + wait_completed = true; + }); + + wait_started_future.wait(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + EXPECT_FALSE(wait_completed); + + // Free up space + queue.Get(); + queue.Get(); + + waiter.join(); + EXPECT_TRUE(wait_completed); +} + +TEST_F(QueueTest, WaitForRoomAtMost) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + // Fill queue partially + producer.Put(1); + producer.Put(2); + producer.Put(3); + + // Should wait until room <= 2 (currently room = 2) + queue.WaitForRoomAtMost(2); + EXPECT_EQ(queue.Size(), 3); + + // Test blocking behavior + std::promise wait_started; + std::future wait_started_future = wait_started.get_future(); + std::atomic wait_completed{false}; + + std::thread waiter([&queue, &wait_started, &wait_completed]() { + wait_started.set_value(); + queue.WaitForRoomAtMost(1); // Should block until room <= 1 + wait_completed = true; + }); + + wait_started_future.wait(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + EXPECT_FALSE(wait_completed); + + // Add one more item to make room = 1 + producer.Put(4); + + waiter.join(); + EXPECT_TRUE(wait_completed); + EXPECT_EQ(queue.Size(), 4); +} + +TEST_F(QueueTest, WaitForSizeAtLeast) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + // Test blocking behavior on empty queue + std::promise wait_started; + std::future wait_started_future = wait_started.get_future(); + std::atomic wait_completed{false}; + + std::thread waiter([&queue, &wait_started, &wait_completed]() { + wait_started.set_value(); + queue.WaitForSizeAtLeast(3); // Should block until size >= 3 + wait_completed = true; + }); + + wait_started_future.wait(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + EXPECT_FALSE(wait_completed); + + // Add items + producer.Put(1); + producer.Put(2); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + EXPECT_FALSE(wait_completed); + + producer.Put(3); // Now size = 3 + + waiter.join(); + EXPECT_TRUE(wait_completed); + EXPECT_EQ(queue.Size(), 3); +} + +TEST_F(QueueTest, WaitForSizeAtMost) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + // Initially empty, size <= 3 + queue.WaitForSizeAtMost(3); + EXPECT_EQ(queue.Size(), 0); + + // Fill queue + producer.Put(1); + producer.Put(2); + producer.Put(3); + producer.Put(4); + producer.Put(5); + + // Test blocking behavior + std::promise wait_started; + std::future wait_started_future = wait_started.get_future(); + std::atomic wait_completed{false}; + + std::thread waiter([&queue, &wait_started, &wait_completed]() { + wait_started.set_value(); + queue.WaitForSizeAtMost(2); // Should block until size <= 2 + wait_completed = true; + }); + + wait_started_future.wait(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + EXPECT_FALSE(wait_completed); + + // Remove items + queue.Get(); + queue.Get(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + EXPECT_FALSE(wait_completed); + + queue.Get(); // Now size = 2 + + waiter.join(); + EXPECT_TRUE(wait_completed); + EXPECT_EQ(queue.Size(), 2); +} + +TEST_F(QueueTest, WaitFunctionsEdgeCases) { + Queue queue(3); + auto producer = queue.CreateProducer(); + + // Wait for room = 0 should work when queue is full + producer.Put(1); + producer.Put(2); + producer.Put(3); + queue.WaitForRoomAtMost(0); + EXPECT_EQ(queue.Size(), 3); + + // Wait for size = 0 should work when queue is empty + queue.Get(); + queue.Get(); + queue.Get(); + queue.WaitForSizeAtMost(0); + EXPECT_EQ(queue.Size(), 0); + + // Wait for room >= capacity should always succeed + queue.WaitForRoomAtLeast(3); + EXPECT_EQ(queue.Size(), 0); + + // Wait for size >= 0 should always succeed + queue.WaitForSizeAtLeast(0); + EXPECT_EQ(queue.Size(), 0); +} + +// Tests for gradual large range operations + +TEST_F(QueueTest, BatchPutAtCapacityWorks) { + Queue queue(3); + auto producer = queue.CreateProducer(); + + // Putting exactly capacity worth of items should work + std::vector items = {1, 2, 3}; + producer.Put(absl::Span(items)); + EXPECT_EQ(queue.Size(), 3); +} + +TEST_F(QueueTest, BatchGetAtCapacityWorks) { + Queue queue(3); + auto producer = queue.CreateProducer(); + + producer.Put(1); + producer.Put(2); + producer.Put(3); + + // Getting exactly capacity worth of items should work + auto result = queue.Get(3); + EXPECT_EQ(result.size(), 3); + EXPECT_EQ(result[0], 1); + EXPECT_EQ(result[1], 2); + EXPECT_EQ(result[2], 3); +} + +TEST_F(QueueTest, LargeRangePutGetGradual) { + Queue queue(3); // Small capacity + auto producer = queue.CreateProducer(); + + // Put more items than capacity - should work gradually + std::vector large_items = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}; + + std::thread put_thread([&producer, &large_items]() { + producer.Put(absl::Span(large_items)); + }); + + // Consume items as they become available + std::vector consumed; + for (int i = 0; i < 10; ++i) { + consumed.push_back(queue.Get()); + } + + put_thread.join(); + + // Verify all items were transferred correctly + EXPECT_EQ(consumed.size(), 10); + for (int i = 0; i < 10; ++i) { + EXPECT_EQ(consumed[i], i + 1); + } + EXPECT_EQ(queue.Size(), 0); +} + +TEST_F(QueueTest, LargeRangePutMove) { + Queue> queue(2); // Very small capacity + auto producer = queue.CreateProducer(); + + // Create large batch of move-only items + std::vector> large_items; + for (int i = 1; i <= 5; ++i) { + large_items.push_back(std::make_unique(i)); + } + + std::thread put_thread([&producer, &large_items]() { + producer.Put(absl::Span>(large_items)); + }); + + // Consume items as they become available + std::vector> consumed; + for (int i = 0; i < 5; ++i) { + consumed.push_back(queue.Get()); + } + + put_thread.join(); + + // Verify all items were transferred correctly + EXPECT_EQ(consumed.size(), 5); + for (int i = 0; i < 5; ++i) { + EXPECT_EQ(*consumed[i], i + 1); + } + EXPECT_EQ(queue.Size(), 0); +} + +TEST_F(QueueTest, LargeRangeGetGradual) { + Queue queue(3); // Small capacity + auto producer = queue.CreateProducer(); + + // Start a thread that will gradually produce items + std::thread producer_thread([&producer]() { + for (int i = 1; i <= 10; ++i) { + producer.Put(i); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + }); + + // Get more items than capacity - should work gradually + auto result = queue.Get(10); + + producer_thread.join(); + + // Verify all items were retrieved correctly + EXPECT_EQ(result.size(), 10); + for (int i = 0; i < 10; ++i) { + EXPECT_EQ(result[i], i + 1); + } + EXPECT_EQ(queue.Size(), 0); +} + +TEST_F(QueueTest, LargeRangePutGetConcurrent) { + Queue queue(5); // Medium capacity + + constexpr int total_items = 100; + constexpr int batch_size = 25; + + auto producer1 = queue.CreateProducer(); + auto producer2 = queue.CreateProducer(); + + std::vector batch1, batch2; + for (int i = 0; i < batch_size; ++i) { + batch1.push_back(i); + batch2.push_back(i + batch_size); + } + + std::atomic items_consumed{0}; + std::vector all_consumed; + all_consumed.reserve(total_items); + + // Multiple producers + std::thread producer_thread1([&producer1, &batch1]() { + producer1.Put(absl::Span(batch1)); + producer1.Put(absl::Span(batch1)); // Put twice + }); + + std::thread producer_thread2([&producer2, &batch2]() { + producer2.Put(absl::Span(batch2)); + producer2.Put(absl::Span(batch2)); // Put twice + }); + + // Consumer getting in large batches + std::thread consumer_thread([&queue, &all_consumed, &items_consumed]() { + while (items_consumed < total_items) { + try { + auto batch = queue.Get(std::min(15, total_items - items_consumed)); + for (const auto& item : batch) { + all_consumed.push_back(item); + } + items_consumed += batch.size(); + } catch (const QueueClosedException&) { + break; + } + } + }); + + producer_thread1.join(); + producer_thread2.join(); + + // Close the queue by destroying producers + producer1.Close(); + producer2.Close(); + + consumer_thread.join(); + + EXPECT_EQ(all_consumed.size(), total_items); + EXPECT_EQ(queue.Size(), 0); +} + +TEST_F(QueueTest, GradualOperationsWithQueueClosure) { + Queue queue(2); // Very small capacity + auto producer = queue.CreateProducer(); + + std::vector large_batch = {1, 2, 3, 4, 5}; + std::atomic exception_caught{false}; + + std::thread producer_thread([&producer, &large_batch, &exception_caught]() { + try { + producer.Put(absl::Span(large_batch)); + } catch (const QueueClosedException&) { + exception_caught = true; + } + }); + + // Let producer start putting items + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + + // Consume a couple items + queue.Get(); // Should get 1 + queue.Get(); // Should get 2 + + // Close the queue while producer is still trying to put items + queue.Close(); + + producer_thread.join(); + + EXPECT_TRUE(exception_caught); + // Queue might have some items that were put before closure + // but we can't predict exactly how many due to timing +} + +// Tests for total put count functionality + +TEST_F(QueueTest, GetTotalPutCountBasic) { + Queue queue(5); + EXPECT_EQ(queue.GetTotalPutCount(), 0); + + { + auto producer = queue.CreateProducer(); + producer.Put(1); + EXPECT_EQ(queue.GetTotalPutCount(), 1); + + producer.Put(2); + producer.Put(3); + EXPECT_EQ(queue.GetTotalPutCount(), 3); + } + + // Count should persist after producer destruction + EXPECT_EQ(queue.GetTotalPutCount(), 3); + + // Count should persist after getting items + queue.Get(); + queue.Get(); + EXPECT_EQ(queue.GetTotalPutCount(), 3); +} + +TEST_F(QueueTest, GetTotalPutCountBatch) { + Queue queue(10); + auto producer = queue.CreateProducer(); + + std::vector batch1 = {1, 2, 3}; + std::vector batch2 = {4, 5}; + + producer.Put(absl::Span(batch1)); + EXPECT_EQ(queue.GetTotalPutCount(), 3); + + producer.Put(absl::Span(batch2)); + EXPECT_EQ(queue.GetTotalPutCount(), 5); + + // Single put after batch + producer.Put(6); + EXPECT_EQ(queue.GetTotalPutCount(), 6); +} + +TEST_F(QueueTest, GetTotalPutCountReset) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + producer.Put(1); + producer.Put(2); + producer.Put(3); + EXPECT_EQ(queue.GetTotalPutCount(), 3); + + // Reset and verify return value + EXPECT_EQ(queue.GetTotalPutCount(true), 3); + EXPECT_EQ(queue.GetTotalPutCount(), 0); + + // Add more items + producer.Put(4); + EXPECT_EQ(queue.GetTotalPutCount(), 1); + + // Non-reset call should not affect counter + EXPECT_EQ(queue.GetTotalPutCount(false), 1); + EXPECT_EQ(queue.GetTotalPutCount(), 1); +} + +TEST_F(QueueTest, GetTotalPutCountThreadSafe) { + Queue queue(50); // Large capacity to avoid blocking + constexpr int items_per_thread = 10; + constexpr int num_threads = 2; + + std::vector threads; + std::vector::Producer> producers; + + // Create producers for each thread + for (int t = 0; t < num_threads; ++t) { + producers.push_back(queue.CreateProducer()); + } + + for (int t = 0; t < num_threads; ++t) { + threads.emplace_back([&producers, t]() { + for (int i = 0; i < items_per_thread; ++i) { + producers[t].Put(t * items_per_thread + i); + } + }); + } + + for (auto& thread : threads) { + thread.join(); + } + + EXPECT_EQ(queue.GetTotalPutCount(), items_per_thread * num_threads); +} + +TEST_F(QueueTest, GetTotalPutCountBatchThreadSafe) { + Queue queue(100); // Large capacity to avoid blocking + + std::vector threads; + std::vector::Producer> producers; + + // Create producers for each thread + for (int t = 0; t < 2; ++t) { + producers.push_back(queue.CreateProducer()); + } + + for (int t = 0; t < 2; ++t) { + threads.emplace_back([&producers, t]() { + std::vector batch; + int batch_size = (t + 1) * 5; // 5, 10 items + for (int i = 0; i < batch_size; ++i) { + batch.push_back(t * 100 + i); + } + producers[t].Put(absl::Span(batch)); + }); + } + + for (auto& thread : threads) { + thread.join(); + } + + EXPECT_EQ(queue.GetTotalPutCount(), 15); // 5 + 10 +} + +TEST_F(QueueTest, GetTotalPutCountWithMoveSemantics) { + Queue> queue(5); + auto producer = queue.CreateProducer(); + + // Single move put + auto ptr1 = std::make_unique(42); + producer.Put(std::move(ptr1)); + EXPECT_EQ(queue.GetTotalPutCount(), 1); + + // Batch move put + std::vector> batch; + for (int i = 0; i < 3; ++i) { + batch.push_back(std::make_unique(i)); + } + producer.Put(absl::Span>(batch)); + EXPECT_EQ(queue.GetTotalPutCount(), 4); +} + +TEST_F(QueueTest, GetTotalPutCountEmptyBatch) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + std::vector empty_batch; + producer.Put(absl::Span(empty_batch)); + EXPECT_EQ(queue.GetTotalPutCount(), 0); + + producer.Put(1); + EXPECT_EQ(queue.GetTotalPutCount(), 1); + + producer.Put(absl::Span(empty_batch)); + EXPECT_EQ(queue.GetTotalPutCount(), 1); +} + +// Tests for DROP_NEW overflow behavior + +TEST_F(QueueTest, DropNewBasicBehavior) { + Queue queue(3, Queue::OverflowBehavior::DROP_NEW); + auto producer = queue.CreateProducer(); + + // Fill queue to capacity + producer.Put(1); + producer.Put(2); + producer.Put(3); + EXPECT_EQ(queue.Size(), 3); + EXPECT_EQ(queue.GetTotalPutCount(), 3); + EXPECT_EQ(queue.GetTotalDropCount(), 0); + + // Additional puts should be dropped + producer.Put(4); + producer.Put(5); + EXPECT_EQ(queue.Size(), 3); + EXPECT_EQ(queue.GetTotalPutCount(), 5); + EXPECT_EQ(queue.GetTotalDropCount(), 2); + + // Verify original items are still there + EXPECT_EQ(queue.Get(), 1); + EXPECT_EQ(queue.Get(), 2); + EXPECT_EQ(queue.Get(), 3); +} + +TEST_F(QueueTest, DropNewBatchBehavior) { + Queue queue(3, Queue::OverflowBehavior::DROP_NEW); + auto producer = queue.CreateProducer(); + + // Fill queue partially + producer.Put(1); + EXPECT_EQ(queue.Size(), 1); + + // Try to put more than capacity allows + std::vector large_batch = {2, 3, 4, 5, 6}; + producer.Put(absl::Span(large_batch)); + + // Only first 2 should fit + EXPECT_EQ(queue.Size(), 3); + EXPECT_EQ(queue.GetTotalPutCount(), 6); // 1 + 5 attempted + EXPECT_EQ(queue.GetTotalDropCount(), 3); // 4, 5, 6 were dropped + + // Verify what's in the queue + EXPECT_EQ(queue.Get(), 1); + EXPECT_EQ(queue.Get(), 2); + EXPECT_EQ(queue.Get(), 3); +} + +TEST_F(QueueTest, DropNewThreadSafety) { + Queue queue(5, Queue::OverflowBehavior::DROP_NEW); + constexpr int num_threads = 3; + constexpr int items_per_thread = 10; + + std::vector threads; + std::vector::Producer> producers; + + for (int t = 0; t < num_threads; ++t) { + producers.push_back(queue.CreateProducer()); + } + + for (int t = 0; t < num_threads; ++t) { + threads.emplace_back([&producers, t]() { + for (int i = 0; i < items_per_thread; ++i) { + producers[t].Put(t * items_per_thread + i); + } + }); + } + + for (auto& thread : threads) { + thread.join(); + } + + // Queue should have at most capacity items + EXPECT_LE(queue.Size(), 5); + // All puts should be counted + EXPECT_EQ(queue.GetTotalPutCount(), num_threads * items_per_thread); + // Some items should have been dropped + EXPECT_GT(queue.GetTotalDropCount(), 0); + // Put count = successful puts + drops + EXPECT_EQ(queue.GetTotalPutCount(), queue.Size() + queue.GetTotalDropCount()); +} + +// Tests for KEEP_NEWEST overflow behavior + +TEST_F(QueueTest, KeepNewestBasicBehavior) { + Queue queue(3, Queue::OverflowBehavior::KEEP_NEWEST); + auto producer = queue.CreateProducer(); + + // Fill queue to capacity + producer.Put(1); + producer.Put(2); + producer.Put(3); + EXPECT_EQ(queue.Size(), 3); + EXPECT_EQ(queue.GetTotalPutCount(), 3); + EXPECT_EQ(queue.GetTotalDropCount(), 0); + + // Additional puts should replace oldest items + producer.Put(4); + EXPECT_EQ(queue.Size(), 3); + EXPECT_EQ(queue.GetTotalPutCount(), 4); + EXPECT_EQ(queue.GetTotalDropCount(), 1); + + producer.Put(5); + EXPECT_EQ(queue.Size(), 3); + EXPECT_EQ(queue.GetTotalPutCount(), 5); + EXPECT_EQ(queue.GetTotalDropCount(), 2); + + // Verify newest items are kept (3, 4, 5) + EXPECT_EQ(queue.Get(), 3); + EXPECT_EQ(queue.Get(), 4); + EXPECT_EQ(queue.Get(), 5); +} + +TEST_F(QueueTest, KeepNewestBatchBehavior) { + Queue queue(3, Queue::OverflowBehavior::KEEP_NEWEST); + auto producer = queue.CreateProducer(); + + // Fill queue + producer.Put(1); + producer.Put(2); + producer.Put(3); + + // Put large batch that exceeds capacity + std::vector large_batch = {4, 5, 6, 7, 8}; + producer.Put(absl::Span(large_batch)); + + // Queue should still have capacity items + EXPECT_EQ(queue.Size(), 3); + EXPECT_EQ(queue.GetTotalPutCount(), 8); + EXPECT_EQ(queue.GetTotalDropCount(), 5); // 1, 2, 3, 4, 5 dropped + + // Should have the newest 3 items (6, 7, 8) + EXPECT_EQ(queue.Get(), 6); + EXPECT_EQ(queue.Get(), 7); + EXPECT_EQ(queue.Get(), 8); +} + +TEST_F(QueueTest, KeepNewestLargeBatch) { + Queue queue(2, Queue::OverflowBehavior::KEEP_NEWEST); + auto producer = queue.CreateProducer(); + + // Put batch larger than capacity + std::vector large_batch = {1, 2, 3, 4, 5}; + producer.Put(absl::Span(large_batch)); + + // Should keep only the last 2 items + EXPECT_EQ(queue.Size(), 2); + EXPECT_EQ(queue.GetTotalPutCount(), 5); + EXPECT_EQ(queue.GetTotalDropCount(), 3); + + EXPECT_EQ(queue.Get(), 4); + EXPECT_EQ(queue.Get(), 5); +} + +// Tests for counter functionality + +TEST_F(QueueTest, GetTotalGetCountBasic) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + EXPECT_EQ(queue.GetTotalGetCount(), 0); + + producer.Put(1); + producer.Put(2); + producer.Put(3); + EXPECT_EQ(queue.GetTotalGetCount(), 0); + + queue.Get(); + EXPECT_EQ(queue.GetTotalGetCount(), 1); + + queue.Get(); + queue.Get(); + EXPECT_EQ(queue.GetTotalGetCount(), 3); +} + +TEST_F(QueueTest, GetTotalGetCountBatch) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + producer.Put(1); + producer.Put(2); + producer.Put(3); + producer.Put(4); + producer.Put(5); + + queue.Get(3); + EXPECT_EQ(queue.GetTotalGetCount(), 3); + + queue.Get(); + EXPECT_EQ(queue.GetTotalGetCount(), 4); + + queue.Get(1); + EXPECT_EQ(queue.GetTotalGetCount(), 5); +} + +TEST_F(QueueTest, GetTotalGetCountReset) { + Queue queue(3); + auto producer = queue.CreateProducer(); + + producer.Put(1); + producer.Put(2); + + queue.Get(); + queue.Get(); + EXPECT_EQ(queue.GetTotalGetCount(), 2); + + EXPECT_EQ(queue.GetTotalGetCount(true), 2); + EXPECT_EQ(queue.GetTotalGetCount(), 0); +} + +TEST_F(QueueTest, GetTotalDropCountBasic) { + Queue queue(2, Queue::OverflowBehavior::DROP_NEW); + auto producer = queue.CreateProducer(); + + EXPECT_EQ(queue.GetTotalDropCount(), 0); + + producer.Put(1); + producer.Put(2); + EXPECT_EQ(queue.GetTotalDropCount(), 0); + + producer.Put(3); // Should be dropped + EXPECT_EQ(queue.GetTotalDropCount(), 1); + + producer.Put(4); // Should be dropped + EXPECT_EQ(queue.GetTotalDropCount(), 2); +} + +TEST_F(QueueTest, GetTotalDropCountKeepNewest) { + Queue queue(2, Queue::OverflowBehavior::KEEP_NEWEST); + auto producer = queue.CreateProducer(); + + producer.Put(1); + producer.Put(2); + EXPECT_EQ(queue.GetTotalDropCount(), 0); + + producer.Put(3); // Should drop 1 + EXPECT_EQ(queue.GetTotalDropCount(), 1); + + producer.Put(4); // Should drop 2 + EXPECT_EQ(queue.GetTotalDropCount(), 2); +} + +TEST_F(QueueTest, GetTotalDropCountReset) { + Queue queue(1, Queue::OverflowBehavior::DROP_NEW); + auto producer = queue.CreateProducer(); + + producer.Put(1); + producer.Put(2); // Dropped + producer.Put(3); // Dropped + EXPECT_EQ(queue.GetTotalDropCount(), 2); + + EXPECT_EQ(queue.GetTotalDropCount(true), 2); + EXPECT_EQ(queue.GetTotalDropCount(), 0); +} + +// MaybeGet() tests + +TEST_F(QueueTest, MaybeGetOnEmptyQueue) { + Queue queue(5); + auto result = queue.MaybeGet(); + EXPECT_FALSE(result.has_value()); +} + +TEST_F(QueueTest, MaybeGetOnNonEmptyQueue) { + Queue queue(5); + { + auto producer = queue.CreateProducer(); + producer.Put(42); + } + + auto result = queue.MaybeGet(); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(*result, 42); + EXPECT_EQ(queue.Size(), 0); +} + +TEST_F(QueueTest, MaybeGetMultipleValues) { + Queue queue(5); + { + auto producer = queue.CreateProducer(); + producer.Put(1); + producer.Put(2); + producer.Put(3); + } + + auto result1 = queue.MaybeGet(); + ASSERT_TRUE(result1.has_value()); + EXPECT_EQ(*result1, 1); + + auto result2 = queue.MaybeGet(); + ASSERT_TRUE(result2.has_value()); + EXPECT_EQ(*result2, 2); + + auto result3 = queue.MaybeGet(); + ASSERT_TRUE(result3.has_value()); + EXPECT_EQ(*result3, 3); + + auto result4 = queue.MaybeGet(); + EXPECT_FALSE(result4.has_value()); +} + +TEST_F(QueueTest, MaybeGetWithMoveOnlyType) { + Queue> queue(5); + { + auto producer = queue.CreateProducer(); + producer.Put(std::make_unique(42)); + } + + auto result = queue.MaybeGet(); + ASSERT_TRUE(result.has_value()); + ASSERT_NE(*result, nullptr); + EXPECT_EQ(**result, 42); +} + +TEST_F(QueueTest, MaybeGetUpdatesGetCount) { + Queue queue(5); + { + auto producer = queue.CreateProducer(); + producer.Put(1); + producer.Put(2); + } + + EXPECT_EQ(queue.GetTotalGetCount(), 0); + + queue.MaybeGet(); + EXPECT_EQ(queue.GetTotalGetCount(), 1); + + queue.MaybeGet(); + EXPECT_EQ(queue.GetTotalGetCount(), 2); + + queue.MaybeGet(); // Empty queue. + EXPECT_EQ(queue.GetTotalGetCount(), 2); // Count not incremented. +} + +// Tests for stop_token cancellation + +TEST_F(QueueTest, StopTokenCancelsPut) { + Queue queue(2); + auto producer = queue.CreateProducer(); + + // Fill the queue + producer.Put(1); + producer.Put(2); + + std::stop_source stop_source; + stop_source.request_stop(); + + // Should immediately throw without blocking since token is already stopped + EXPECT_THROW(producer.Put(3, stop_source.get_token()), QueueRequestCancelled); + EXPECT_EQ(queue.Size(), 2); +} + +TEST_F(QueueTest, StopTokenCancelsGet) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + std::stop_source stop_source; + stop_source.request_stop(); + + // Should immediately throw without blocking since token is already stopped + EXPECT_THROW(queue.Get(stop_source.get_token()), QueueRequestCancelled); + EXPECT_EQ(queue.Size(), 0); +} + +TEST_F(QueueTest, StopTokenCancelsBatchPut) { + Queue queue(2); + auto producer = queue.CreateProducer(); + + // Fill queue completely + producer.Put(1); + producer.Put(2); + + std::stop_source stop_source; + stop_source.request_stop(); + + // Should immediately throw without blocking since token is already stopped + std::vector items = {3, 4, 5}; + EXPECT_THROW( + producer.Put(absl::Span(items), stop_source.get_token()), + QueueRequestCancelled); +} + +TEST_F(QueueTest, StopTokenCancelsBatchGet) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + std::stop_source stop_source; + stop_source.request_stop(); + + // Should immediately throw without blocking since token is already stopped + EXPECT_THROW(queue.Get(10, stop_source.get_token()), QueueRequestCancelled); +} + +TEST_F(QueueTest, StopTokenCancelsWaitForRoomAtLeast) { + Queue queue(3); + auto producer = queue.CreateProducer(); + + // Fill queue + producer.Put(1); + producer.Put(2); + producer.Put(3); + + std::stop_source stop_source; + stop_source.request_stop(); + + // Should immediately throw without blocking since token is already stopped + EXPECT_THROW(queue.WaitForRoomAtLeast(2, stop_source.get_token()), + QueueRequestCancelled); +} + +TEST_F(QueueTest, StopTokenCancelsWaitForSizeAtLeast) { + Queue queue(5); + auto producer = queue.CreateProducer(); + + std::stop_source stop_source; + stop_source.request_stop(); + + // Should immediately throw without blocking since token is already stopped + EXPECT_THROW(queue.WaitForSizeAtLeast(3, stop_source.get_token()), + QueueRequestCancelled); +} + +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/stream_shuffler.cc b/csrc/utils/stream_shuffler.cc new file mode 100644 index 00000000..d1acc41f --- /dev/null +++ b/csrc/utils/stream_shuffler.cc @@ -0,0 +1,109 @@ +#include "utils/stream_shuffler.h" + +namespace lczero { +namespace training { + +void StreamShuffler::SetUpperBound(size_t upper_bound) { + assert(upper_bound >= upper_bound_); + stream_size_ += upper_bound - upper_bound_; + while (upper_bound_ < upper_bound) { + if (buckets_.empty() || buckets_.back().GetRemainingCapacity() == 0) { + buckets_.emplace_back(upper_bound_, bucket_size_); + } + upper_bound_ = std::min( + upper_bound, upper_bound_ + buckets_.back().GetRemainingCapacity()); + buckets_.back().Extend(upper_bound_); + } +} + +void StreamShuffler::SetLowerBound(size_t lower_bound) { + assert(lower_bound >= lower_bound_); + lower_bound_ = lower_bound; + if (lower_bound >= upper_bound_) { + upper_bound_ = lower_bound; + stream_size_ = 0; + buckets_.clear(); + return; + } + while (!buckets_.empty() && buckets_.front().upper_bound() <= lower_bound_) { + stream_size_ -= buckets_.front().size(); + buckets_.pop_front(); + } + if (!buckets_.empty()) { + auto old_size = buckets_.front().size(); + buckets_.front().DeclareLowerBound(lower_bound_); + stream_size_ -= old_size - buckets_.front().size(); + } +} + +std::optional StreamShuffler::GetNextItem() { + auto try_fetch = [&]() -> size_t { + size_t item_idx = absl::Uniform(gen_, size_t{0}, stream_size_); + --stream_size_; + for (auto& bucket : buckets_) { + if (item_idx < bucket.size()) return bucket.Fetch(item_idx); + item_idx -= bucket.size(); + } + throw std::logic_error("StreamShuffler: item index out of bounds"); + }; + + while (stream_size_ > 0) { + if (auto item = try_fetch(); item >= lower_bound_) return item; + } + + return std::nullopt; +} + +void StreamShuffler::Reset(size_t lower_bound, size_t upper_bound) { + // Reset all internal state + buckets_.clear(); + stream_size_ = 0; + upper_bound_ = lower_bound; + lower_bound_ = lower_bound; + + // Establish the bounds, which will build the buckets with fresh data + if (upper_bound > lower_bound) { + SetUpperBound(upper_bound); + } +} + +StreamShuffler::Bucket::Bucket(size_t lower_bound, size_t capacity) + : upper_bound_(lower_bound), items_(capacity) {} + +size_t StreamShuffler::Bucket::GetRemainingCapacity() const { + return items_.size() - items_count_; +} + +void StreamShuffler::Bucket::Extend(size_t new_upper_bound) { + assert(new_upper_bound >= upper_bound_); + const size_t increase = new_upper_bound - upper_bound_; + assert(increase <= GetRemainingCapacity()); + std::iota(items_.begin() + items_count_, + items_.begin() + items_count_ + increase, upper_bound_); + items_count_ += increase; + upper_bound_ = new_upper_bound; +} + +size_t StreamShuffler::Bucket::Fetch(size_t item_idx) { + assert(item_idx < items_count_); + size_t item = items_[item_idx]; + std::swap(items_[item_idx], items_[--items_count_]); + return item; +} + +void DeclareLowerBound(size_t new_lower_bound); +void StreamShuffler::Bucket::DeclareLowerBound(size_t new_lower_bound) { + if (upper_bound_ - new_lower_bound < 2 * items_count_) return; + + // If the bucket has much more items that the allowed range, there are many + // items out of the range. It makes sense to sort and remove them. + std::sort(items_.begin(), items_.begin() + items_count_, + std::greater()); + // Find the first item that is under the new lower bound. + auto it = std::upper_bound(items_.begin(), items_.begin() + items_count_, + new_lower_bound, std::greater()); + items_count_ = it - items_.begin(); +} + +} // namespace training +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/stream_shuffler.h b/csrc/utils/stream_shuffler.h new file mode 100644 index 00000000..86007b96 --- /dev/null +++ b/csrc/utils/stream_shuffler.h @@ -0,0 +1,61 @@ +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace lczero { +namespace training { + +// Returns a number between [lower_bound, upper_bound) in shuffled order. +// Both bounds can be changed at any time, and the stream will adapt +// accordingly. Not thread-safe. +class StreamShuffler { + public: + // Sets the upper bound (exclusive). Can only be increased. + void SetUpperBound(size_t upper_bound); + + // Sets the lower bound (inclusive). Can only be increased. + void SetLowerBound(size_t lower_bound); + + // Sets the bucket size for internal storage optimization. + void SetBucketSize(size_t bucket_size) { bucket_size_ = bucket_size; } + + // Returns the next item in shuffled order, or nullopt if exhausted. + std::optional GetNextItem(); + + // Resets the shuffler to restart iteration with specified bounds. + void Reset(size_t lower_bound, size_t upper_bound); + + private: + class Bucket { + public: + Bucket(size_t lower_bound, size_t capacity); + size_t GetRemainingCapacity() const; + void Extend(size_t new_upper_bound); + size_t Fetch(size_t item_idx); + void DeclareLowerBound(size_t new_lower_bound); + + size_t upper_bound() const { return upper_bound_; } + size_t size() const { return items_count_; } + + private: + size_t upper_bound_ = 0; + size_t items_count_ = 0; + absl::FixedArray items_; + }; + + absl::BitGen gen_; + std::deque buckets_; + size_t stream_size_ = 0; + size_t upper_bound_ = 0; + size_t lower_bound_ = 0; + size_t bucket_size_ = 524288; +}; + +} // namespace training +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/stream_shuffler_test.cc b/csrc/utils/stream_shuffler_test.cc new file mode 100644 index 00000000..f606a218 --- /dev/null +++ b/csrc/utils/stream_shuffler_test.cc @@ -0,0 +1,287 @@ +#include "utils/stream_shuffler.h" + +#include +#include +#include + +#include +#include + +namespace lczero { +namespace training { + +class StreamShufflerTest : public ::testing::Test { + protected: + void SetUp() override { shuffler_.SetBucketSize(4); } + + StreamShuffler shuffler_; +}; + +TEST_F(StreamShufflerTest, EmptyRangeReturnsNullopt) { + shuffler_.SetUpperBound(10); + shuffler_.SetLowerBound(10); + EXPECT_EQ(shuffler_.GetNextItem(), std::nullopt); +} + +TEST_F(StreamShufflerTest, SingleItemRange) { + shuffler_.SetUpperBound(1); + shuffler_.SetLowerBound(0); + + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + EXPECT_EQ(item.value(), 0); + + EXPECT_EQ(shuffler_.GetNextItem(), std::nullopt); +} + +TEST_F(StreamShufflerTest, BasicRangeGeneration) { + shuffler_.SetUpperBound(5); + shuffler_.SetLowerBound(0); + + std::set received; + for (int i = 0; i < 5; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + EXPECT_GE(item.value(), 0); + EXPECT_LT(item.value(), 5); + EXPECT_TRUE(received.insert(item.value()).second); + } + + EXPECT_EQ(received.size(), 5); + EXPECT_EQ(shuffler_.GetNextItem(), std::nullopt); +} + +TEST_F(StreamShufflerTest, HeadAdvancesByBucketMultiples) { + shuffler_.SetUpperBound(4); + shuffler_.SetLowerBound(0); + + std::set received; + for (int i = 0; i < 4; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + received.insert(item.value()); + } + EXPECT_EQ(received.size(), 4); + + shuffler_.SetUpperBound(8); + for (int i = 0; i < 4; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + EXPECT_GE(item.value(), 0); + EXPECT_LT(item.value(), 8); + EXPECT_TRUE(received.insert(item.value()).second); + } + EXPECT_EQ(received.size(), 8); + EXPECT_EQ(shuffler_.GetNextItem(), std::nullopt); +} + +TEST_F(StreamShufflerTest, HeadAdvancesByNonMultiples) { + shuffler_.SetUpperBound(3); + shuffler_.SetLowerBound(0); + + std::set received; + for (int i = 0; i < 3; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + received.insert(item.value()); + } + + shuffler_.SetUpperBound(7); + for (int i = 0; i < 4; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + EXPECT_GE(item.value(), 0); + EXPECT_LT(item.value(), 7); + EXPECT_TRUE(received.insert(item.value()).second); + } + EXPECT_EQ(received.size(), 7); + EXPECT_EQ(shuffler_.GetNextItem(), std::nullopt); +} + +TEST_F(StreamShufflerTest, TailAdvancesByBucketMultiples) { + shuffler_.SetUpperBound(12); + shuffler_.SetLowerBound(0); + + std::set all_received; + for (int i = 0; i < 4; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + EXPECT_GE(item.value(), 0); + EXPECT_LT(item.value(), 12); + EXPECT_TRUE(all_received.insert(item.value()).second); + } + + shuffler_.SetLowerBound(4); + std::optional item; + while ((item = shuffler_.GetNextItem()).has_value()) { + EXPECT_GE(item.value(), 4); + EXPECT_LT(item.value(), 12); + EXPECT_TRUE(all_received.insert(item.value()).second); + } + + // Verify all items in range [4, 12) were eventually fetched + for (size_t i = 4; i < 12; ++i) { + EXPECT_TRUE(all_received.count(i) > 0) + << "Item " << i << " was never fetched"; + } +} + +TEST_F(StreamShufflerTest, TailAdvancesByNonMultiples) { + shuffler_.SetUpperBound(10); + shuffler_.SetLowerBound(0); + + std::set all_received; + for (int i = 0; i < 3; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + EXPECT_GE(item.value(), 0); + EXPECT_LT(item.value(), 10); + EXPECT_TRUE(all_received.insert(item.value()).second); + } + + shuffler_.SetLowerBound(3); + std::optional item; + while ((item = shuffler_.GetNextItem()).has_value()) { + EXPECT_GE(item.value(), 3); + EXPECT_LT(item.value(), 10); + EXPECT_TRUE(all_received.insert(item.value()).second); + } + + // Verify all items in range [3, 10) were eventually fetched + for (size_t i = 3; i < 10; ++i) { + EXPECT_TRUE(all_received.count(i) > 0) + << "Item " << i << " was never fetched"; + } +} + +TEST_F(StreamShufflerTest, BothBoundsSlideSimultaneously) { + shuffler_.SetUpperBound(10); + shuffler_.SetLowerBound(0); + + std::set all_received; + for (int i = 0; i < 5; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + EXPECT_GE(item.value(), 0); + EXPECT_LT(item.value(), 10); + EXPECT_TRUE(all_received.insert(item.value()).second); + } + + shuffler_.SetUpperBound(15); + shuffler_.SetLowerBound(5); + + std::optional item; + while ((item = shuffler_.GetNextItem()).has_value()) { + EXPECT_GE(item.value(), 5); + EXPECT_LT(item.value(), 15); + EXPECT_TRUE(all_received.insert(item.value()).second); + } + + // Verify all items in range [5, 15) were eventually fetched + for (size_t i = 5; i < 15; ++i) { + EXPECT_TRUE(all_received.count(i) > 0) + << "Item " << i << " was never fetched"; + } +} + +TEST_F(StreamShufflerTest, ComplexSlidingWindow) { + std::set all_received; + + shuffler_.SetUpperBound(6); + shuffler_.SetLowerBound(0); + + for (int i = 0; i < 3; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + all_received.insert(item.value()); + } + + shuffler_.SetUpperBound(11); + for (int i = 0; i < 2; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + all_received.insert(item.value()); + } + + shuffler_.SetLowerBound(2); + shuffler_.SetUpperBound(14); + + std::set final_received; + std::optional item; + while ((item = shuffler_.GetNextItem()).has_value()) { + EXPECT_GE(item.value(), 2); + EXPECT_LT(item.value(), 14); + EXPECT_TRUE(final_received.insert(item.value()).second); + } + + for (const auto& val : final_received) { + EXPECT_GE(val, 2); + EXPECT_LT(val, 14); + } +} + +TEST_F(StreamShufflerTest, UniquenessAcrossMultipleBuckets) { + shuffler_.SetUpperBound(20); + shuffler_.SetLowerBound(0); + + std::set received; + std::optional item; + while ((item = shuffler_.GetNextItem()).has_value()) { + EXPECT_GE(item.value(), 0); + EXPECT_LT(item.value(), 20); + EXPECT_TRUE(received.insert(item.value()).second); + } + + EXPECT_EQ(received.size(), 20); +} + +TEST_F(StreamShufflerTest, TailCatchesUpToHead) { + shuffler_.SetUpperBound(8); + shuffler_.SetLowerBound(0); + + for (int i = 0; i < 3; ++i) { + auto item = shuffler_.GetNextItem(); + ASSERT_TRUE(item.has_value()); + } + + shuffler_.SetLowerBound(8); + EXPECT_EQ(shuffler_.GetNextItem(), std::nullopt); +} + +TEST_F(StreamShufflerTest, ResetAllowsIterationRestart) { + shuffler_.SetUpperBound(5); + shuffler_.SetLowerBound(0); + + // Exhaust all items + absl::flat_hash_set first_round; + std::optional item; + while ((item = shuffler_.GetNextItem()).has_value()) { + first_round.insert(item.value()); + } + + // Should have gotten all 5 items + EXPECT_EQ(first_round.size(), 5); + + // Shuffler should be exhausted + EXPECT_EQ(shuffler_.GetNextItem(), std::nullopt); + + // Reset the shuffler + shuffler_.Reset(0, 5); + + // Should be able to get items again + absl::flat_hash_set second_round; + int count = 0; + while ((item = shuffler_.GetNextItem()).has_value() && count < 10) { + second_round.insert(item.value()); + count++; + } + + // Should get all items again + EXPECT_EQ(second_round.size(), 5); + + // Both rounds should contain the same set of items + EXPECT_EQ(first_round, second_round); +} + +} // namespace training +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/tensor.h b/csrc/utils/tensor.h new file mode 100644 index 00000000..a4fa6d72 --- /dev/null +++ b/csrc/utils/tensor.h @@ -0,0 +1,121 @@ +#pragma once + +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/container/fixed_array.h" +#include "absl/types/span.h" + +namespace lczero { + +// Class that holds tensor which will be exposed through pybind11. +class TensorBase { + public: + virtual ~TensorBase() = default; + virtual void* data() = 0; + virtual const void* data() const = 0; + virtual const std::vector& shape() const = 0; + virtual const std::vector& strides() const = 0; + virtual size_t element_size() const = 0; + virtual std::string py_format() const = 0; +}; + +template +class TypedTensor : public TensorBase { + public: + TypedTensor(std::initializer_list shape) + : data_(CalculateTotalSize(shape)), shape_(shape.begin(), shape.end()) { + // Calculate strides in row-major order (in bytes). + strides_.resize(shape_.size()); + size_t total_size = 1; + for (int i = static_cast(shape_.size()) - 1; i >= 0; --i) { + strides_[i] = total_size * sizeof(T); + total_size *= shape_[i]; + } + } + + void* data() override { return data_.data(); } + + const void* data() const override { return data_.data(); } + + const std::vector& shape() const override { return shape_; } + + const std::vector& strides() const override { return strides_; } + + size_t element_size() const override { return sizeof(T); } + + std::string py_format() const override { + if constexpr (std::is_same_v) { + return "f"; + } else if constexpr (std::is_same_v) { + return "d"; + } else if constexpr (std::is_same_v) { + return "i"; + } else if constexpr (std::is_same_v) { + return "q"; + } else { + static_assert(std::is_same_v, "Unsupported tensor type"); + } + } + + T& operator[](absl::Span dims) { + if (dims.size() != shape_.size()) { + throw std::invalid_argument( + "Number of dimensions must match tensor rank"); + } + return data_[CalculateOffset(dims)]; + } + + const T& operator[](absl::Span dims) const { + if (dims.size() != shape_.size()) { + throw std::invalid_argument( + "Number of dimensions must match tensor rank"); + } + return data_[CalculateOffset(dims)]; + } + + absl::Span slice(absl::Span dims) { + if (dims.size() > shape_.size()) { + throw std::invalid_argument( + "Number of dimensions cannot exceed tensor rank"); + } + return absl::Span(data_.data() + CalculateOffset(dims), + CalculateSliceSize(dims.size())); + } + + absl::Span slice(absl::Span dims) const { + if (dims.size() > shape_.size()) { + throw std::invalid_argument( + "Number of dimensions cannot exceed tensor rank"); + } + return absl::Span(data_.data() + CalculateOffset(dims), + CalculateSliceSize(dims.size())); + } + + private: + size_t CalculateOffset(absl::Span dims) const { + return absl::c_inner_product( + dims, strides_, size_t{0}, std::plus<>{}, + [](ssize_t dim, ssize_t stride) { return dim * stride / sizeof(T); }); + } + + size_t CalculateSliceSize(size_t dims_size) const { + return absl::c_accumulate(absl::MakeConstSpan(shape_).subspan(dims_size), + size_t{1}, std::multiplies{}); + } + + static size_t CalculateTotalSize(std::initializer_list shape) { + return absl::c_accumulate(shape, size_t{1}, std::multiplies{}); + } + + absl::FixedArray data_; + std::vector shape_; + std::vector strides_; +}; + +using TensorTuple = std::vector>; + +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/tensor_test.cc b/csrc/utils/tensor_test.cc new file mode 100644 index 00000000..62a77587 --- /dev/null +++ b/csrc/utils/tensor_test.cc @@ -0,0 +1,160 @@ +// ABOUTME: Unit tests for tensor classes and their data access methods. +// ABOUTME: Tests construction, element access, slicing, and error conditions. + +#include "utils/tensor.h" + +#include + +namespace lczero { +namespace { + +TEST(TypedTensorTest, ConstructorAndBasicProperties) { + TypedTensor tensor({2, 3, 4}); + + // Check shape + EXPECT_EQ(tensor.shape().size(), 3); + EXPECT_EQ(tensor.shape()[0], 2); + EXPECT_EQ(tensor.shape()[1], 3); + EXPECT_EQ(tensor.shape()[2], 4); + + // Check strides (in bytes, row-major order) + EXPECT_EQ(tensor.strides().size(), 3); + EXPECT_EQ(tensor.strides()[0], 12 * sizeof(float)); // 3 * 4 elements + EXPECT_EQ(tensor.strides()[1], 4 * sizeof(float)); // 4 elements + EXPECT_EQ(tensor.strides()[2], 1 * sizeof(float)); // 1 element + + // Check element size + EXPECT_EQ(tensor.element_size(), sizeof(float)); + + // Check py_format + EXPECT_EQ(tensor.py_format(), "f"); + + // Check data pointer is valid + EXPECT_NE(tensor.data(), nullptr); +} + +TEST(TypedTensorTest, PyFormatForDifferentTypes) { + TypedTensor float_tensor({2}); + EXPECT_EQ(float_tensor.py_format(), "f"); + + TypedTensor double_tensor({2}); + EXPECT_EQ(double_tensor.py_format(), "d"); + + TypedTensor int32_tensor({2}); + EXPECT_EQ(int32_tensor.py_format(), "i"); + + TypedTensor int64_tensor({2}); + EXPECT_EQ(int64_tensor.py_format(), "q"); +} + +TEST(TypedTensorTest, ElementAccess) { + TypedTensor tensor({2, 3}); + + // Set some values + tensor[{0, 0}] = 10; + tensor[{0, 1}] = 11; + tensor[{0, 2}] = 12; + tensor[{1, 0}] = 20; + tensor[{1, 1}] = 21; + tensor[{1, 2}] = 22; + + // Check values + EXPECT_EQ((tensor[{0, 0}]), 10); + EXPECT_EQ((tensor[{0, 1}]), 11); + EXPECT_EQ((tensor[{0, 2}]), 12); + EXPECT_EQ((tensor[{1, 0}]), 20); + EXPECT_EQ((tensor[{1, 1}]), 21); + EXPECT_EQ((tensor[{1, 2}]), 22); +} + +TEST(TypedTensorTest, ConstElementAccess) { + TypedTensor tensor({2, 2}); + tensor[{0, 0}] = 1; + tensor[{0, 1}] = 2; + tensor[{1, 0}] = 3; + tensor[{1, 1}] = 4; + + const auto& const_tensor = tensor; + EXPECT_EQ((const_tensor[{0, 0}]), 1); + EXPECT_EQ((const_tensor[{0, 1}]), 2); + EXPECT_EQ((const_tensor[{1, 0}]), 3); + EXPECT_EQ((const_tensor[{1, 1}]), 4); +} + +TEST(TypedTensorTest, SliceAccess) { + TypedTensor tensor({2, 3, 4}); + + // Fill with test data + for (int i = 0; i < 2; ++i) { + for (int j = 0; j < 3; ++j) { + for (int k = 0; k < 4; ++k) { + tensor[{i, j, k}] = i * 100 + j * 10 + k; + } + } + } + + // Test 1D slice (fix first dimension) + auto slice1d = tensor.slice({1}); + EXPECT_EQ(slice1d.size(), 12); // 3 * 4 elements + EXPECT_EQ(slice1d[0], 100); // tensor[{1, 0, 0}] + EXPECT_EQ(slice1d[4], 110); // tensor[{1, 1, 0}] + + // Test 2D slice (fix first two dimensions) + auto slice2d = tensor.slice({0, 1}); + EXPECT_EQ(slice2d.size(), 4); // 4 elements + EXPECT_EQ(slice2d[0], 10); // tensor[{0, 1, 0}] + EXPECT_EQ(slice2d[1], 11); // tensor[{0, 1, 1}] + EXPECT_EQ(slice2d[2], 12); // tensor[{0, 1, 2}] + EXPECT_EQ(slice2d[3], 13); // tensor[{0, 1, 3}] + + // Test full tensor slice (no dimensions fixed) + auto full_slice = tensor.slice({}); + EXPECT_EQ(full_slice.size(), 24); // 2 * 3 * 4 elements +} + +TEST(TypedTensorTest, ConstSliceAccess) { + TypedTensor tensor({2, 2}); + tensor[{0, 0}] = 1; + tensor[{0, 1}] = 2; + tensor[{1, 0}] = 3; + tensor[{1, 1}] = 4; + + const auto& const_tensor = tensor; + auto slice = const_tensor.slice({0}); + EXPECT_EQ(slice.size(), 2); + EXPECT_EQ(slice[0], 1); + EXPECT_EQ(slice[1], 2); +} + +TEST(TypedTensorTest, ElementAccessWrongDimensions) { + TypedTensor tensor({2, 3}); + + EXPECT_THROW((tensor[{0}]), std::invalid_argument); + EXPECT_THROW((tensor[{0, 1, 2}]), std::invalid_argument); +} + +TEST(TypedTensorTest, SliceAccessTooManyDimensions) { + TypedTensor tensor({2, 3}); + + EXPECT_THROW((tensor.slice({0, 1, 2})), std::invalid_argument); +} + +TEST(TypedTensorTest, OneDimensionalTensor) { + TypedTensor tensor({5}); + + EXPECT_EQ(tensor.shape().size(), 1); + EXPECT_EQ(tensor.shape()[0], 5); + EXPECT_EQ(tensor.strides()[0], sizeof(float)); + + tensor[{0}] = 1.0f; + tensor[{4}] = 5.0f; + + EXPECT_EQ((tensor[{0}]), 1.0f); + EXPECT_EQ((tensor[{4}]), 5.0f); + + auto slice = tensor.slice({}); + EXPECT_EQ(slice.size(), 5); +} + +} // namespace +} // namespace lczero \ No newline at end of file diff --git a/csrc/utils/thread_pool.h b/csrc/utils/thread_pool.h new file mode 100644 index 00000000..6ee059dd --- /dev/null +++ b/csrc/utils/thread_pool.h @@ -0,0 +1,228 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/functional/any_invocable.h" +#include "absl/synchronization/mutex.h" + +namespace lczero { + +struct ThreadPoolOptions { + // If true, starts new thread when task is enqueued and no threads are idle. + bool grow_automatically = false; +}; + +class ThreadPool { + public: + ThreadPool(size_t initial_threads = 0, + const ThreadPoolOptions& options = ThreadPoolOptions(), + std::stop_source stop_source = std::stop_source()); + + // Blocks until all tasks are completed and threads are joined. + ~ThreadPool(); + + // Returns the stop_token for this thread pool. + std::stop_token stop_token() const; + + // Enqueues a task for execution and returns a std::future. + // If the provided function accepts a std::stop_token as its first argument, + // one will be passed to it from the thread pool's stop source. + template + auto Enqueue(F&& f, Args&&... args) { + // This lambda captures the common queuing logic. + auto enqueue_common = [&](auto&& task_to_enqueue, auto&& future_to_return) { + { + absl::MutexLock lock(&mutex_); + running_tasks_ += 1; + while (options_.grow_automatically && + running_tasks_ >= threads_.size()) { + StartWorkerThread(); + } + pending_tasks_.emplace_back( + [task = std::move(task_to_enqueue)]() mutable { task(); }); + work_available_.Signal(); + } + return future_to_return; + }; + + if constexpr (std::is_invocable_v) { + using ReturnType = std::invoke_result_t; + std::packaged_task task( + std::bind(std::forward(f), stop_source_.get_token(), + std::forward(args)...)); + std::future future = task.get_future(); + return enqueue_common(std::move(task), std::move(future)); + } else { + using ReturnType = std::invoke_result_t; + std::packaged_task task( + std::bind(std::forward(f), std::forward(args)...)); + std::future future = task.get_future(); + return enqueue_common(std::move(task), std::move(future)); + } + } + + // Waits for all tasks to complete. + void WaitAll(); + + // Waits for at least one thread to become available, i.e. no tasks are + // pending, and number of running tasks is less than the number of threads. + void WaitForAvailableThread(); + + // Waits until the number of queued but not yet started tasks is below + // the specified threshold. + void WaitForPendingTasksBelow(size_t threshold); + + // Number of tasks that are not yet started. + size_t num_pending_tasks() const; + + // Number of tasks that are currently running. + size_t num_running_tasks() const; + + // Number of worker threads (busy or not). + size_t num_threads() const; + + // Signal workers to terminate and join all threads. + void Shutdown(); + + private: + void WorkerLoop(); + void WorkerEntryPoint(); + void StartWorkerThread() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + bool AllTasksCompletedCond() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + return pending_tasks_.empty() && running_tasks_ == 0; + } + bool ThreadAvailableCond() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + return pending_tasks_.empty() && running_tasks_ < threads_.size(); + } + + ThreadPool(const ThreadPool&) = delete; + ThreadPool& operator=(const ThreadPool&) = delete; + ThreadPool(ThreadPool&&) = delete; + ThreadPool& operator=(ThreadPool&&) = delete; + + ThreadPoolOptions options_; + mutable absl::Mutex mutex_; + absl::CondVar work_available_; + absl::CondVar work_done_; + + std::stop_source stop_source_; + std::vector threads_ ABSL_GUARDED_BY(mutex_); + std::deque> pending_tasks_ ABSL_GUARDED_BY(mutex_); + size_t running_tasks_ ABSL_GUARDED_BY(mutex_) = 0; +}; + +inline ThreadPool::ThreadPool(size_t initial_threads, + const ThreadPoolOptions& options, + std::stop_source stop_source) + : options_(options), stop_source_(std::move(stop_source)) { + absl::MutexLock lock(&mutex_); + for (size_t i = 0; i < initial_threads; ++i) { + StartWorkerThread(); + } +} + +inline ThreadPool::~ThreadPool() { Shutdown(); } + +inline std::stop_token ThreadPool::stop_token() const { + return stop_source_.get_token(); +} + +inline void ThreadPool::WorkerLoop() { + while (true) { + absl::AnyInvocable task; + { + absl::MutexLock lock(&mutex_); + while (!stop_source_.stop_requested() && pending_tasks_.empty()) { + work_available_.Wait(&mutex_); + } + if (stop_source_.stop_requested() && pending_tasks_.empty()) return; + task = std::move(pending_tasks_.front()); + pending_tasks_.pop_front(); + } + + std::move(task)(); + + { + absl::MutexLock lock(&mutex_); + running_tasks_ -= 1; + work_done_.SignalAll(); + if (!pending_tasks_.empty()) work_available_.Signal(); + } + } +} + +inline void ThreadPool::WorkerEntryPoint() { + try { + WorkerLoop(); + } catch (const std::exception& exception) { + std::cerr << "ThreadPool worker exited due to uncaught exception: " + << exception.what() << std::endl; + throw; + } catch (...) { + std::cerr << "ThreadPool worker exited due to unknown exception." + << std::endl; + throw; + } +} + +inline void ThreadPool::WaitAll() { + absl::MutexLock lock(&mutex_); + while (!AllTasksCompletedCond()) { + work_done_.Wait(&mutex_); + } +} + +inline void ThreadPool::WaitForAvailableThread() { + absl::MutexLock lock(&mutex_); + while (!ThreadAvailableCond()) { + work_done_.Wait(&mutex_); + } +} + +inline void ThreadPool::WaitForPendingTasksBelow(size_t threshold) { + absl::MutexLock lock(&mutex_); + while (pending_tasks_.size() >= threshold) { + work_done_.Wait(&mutex_); + } +} + +inline void ThreadPool::StartWorkerThread() + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + threads_.emplace_back(&ThreadPool::WorkerEntryPoint, this); +} + +inline size_t ThreadPool::num_pending_tasks() const { + absl::MutexLock lock(&mutex_); + return pending_tasks_.size(); +} + +inline size_t ThreadPool::num_running_tasks() const { + absl::MutexLock lock(&mutex_); + return std::max(running_tasks_, threads_.size()); +} + +inline size_t ThreadPool::num_threads() const { + absl::MutexLock lock(&mutex_); + return threads_.size(); +} + +inline void ThreadPool::Shutdown() { + { + absl::MutexLock lock(&mutex_); + if (!stop_source_.stop_requested()) { + stop_source_.request_stop(); + } + work_available_.SignalAll(); + work_done_.SignalAll(); + } + threads_.clear(); +} + +} // namespace lczero diff --git a/csrc/utils/training_data_printer.cc b/csrc/utils/training_data_printer.cc new file mode 100644 index 00000000..c87c4dd4 --- /dev/null +++ b/csrc/utils/training_data_printer.cc @@ -0,0 +1,122 @@ +#include "utils/training_data_printer.h" + +#include + +#include +#include + +#include "chess/board.h" +#include "neural/decoder.h" +#include "trainingdata/reader.h" + +namespace lczero { +namespace training { + +void PrintFloatArray(const float* data, size_t size, absl::string_view name, + int64_t per_line) { + per_line = std::max(1, per_line); + std::cout << " " << name << ":\n"; + for (size_t i = 0; i < size; ++i) { + if (i % per_line == 0) { + std::cout << " [" << absl::StrFormat("%4zu", i) << "]: "; + } + std::cout << absl::StrFormat("% .6g", data[i]); + if ((i + 1) % per_line == 0 || i + 1 == size) { + std::cout << "\n"; + } else { + std::cout << ", "; + } + } +} + +void PrintUint64Array(const uint64_t* data, size_t size, absl::string_view name, + int64_t per_line) { + per_line = std::max(1, per_line); + std::cout << " " << name << ":\n"; + for (size_t i = 0; i < size; ++i) { + if (i % per_line == 0) { + std::cout << " [" << absl::StrFormat("%3zu", i) << "]: "; + } + std::cout << absl::StrFormat("0x%016x", data[i]); + if ((i + 1) % per_line == 0 || i + 1 == size) { + std::cout << "\n"; + } else { + std::cout << ", "; + } + } +} + +std::string DecodeInvarianceInfo(uint8_t invariance_info) { + return absl::StrFormat( + "flip=%d, mirror=%d, transpose=%d, best_move_proven=%d, " + "max_length=%d, adjudicated=%d, rescorer_deleted=%d, side_to_move=%d", + invariance_info & 0x1, (invariance_info >> 1) & 0x1, + (invariance_info >> 2) & 0x1, (invariance_info >> 3) & 0x1, + (invariance_info >> 4) & 0x1, (invariance_info >> 5) & 0x1, + (invariance_info >> 6) & 0x1, (invariance_info >> 7) & 0x1); +} + +std::string TrainingDataToFen(const FrameType& entry) { + InputPlanes planes = PlanesFromTrainingData(entry); + ChessBoard board; + int rule50 = 0; + int gameply = 0; + PopulateBoard( + static_cast(entry.input_format), + planes, &board, &rule50, &gameply); + std::string fen = BoardToFen(board); + fen += " " + std::to_string(rule50); + fen += " " + std::to_string((gameply / 2) + 1); + return fen; +} + +void PrintTrainingDataEntry(const FrameType& entry, + absl::string_view header_text, + int64_t float_per_line, int64_t plane_per_line) { + std::cout << header_text << "\n"; + std::cout << " FEN: " << TrainingDataToFen(entry) << "\n"; + std::cout << " version: " << entry.version << "\n"; + std::cout << " input_format: " << entry.input_format << "\n"; + std::cout << " castling_us_ooo: " << static_cast(entry.castling_us_ooo) + << "\n"; + std::cout << " castling_us_oo: " << static_cast(entry.castling_us_oo) + << "\n"; + std::cout << " castling_them_ooo: " + << static_cast(entry.castling_them_ooo) << "\n"; + std::cout << " castling_them_oo: " + << static_cast(entry.castling_them_oo) << "\n"; + std::cout << " side_to_move_or_enpassant: " + << static_cast(entry.side_to_move_or_enpassant) << "\n"; + std::cout << " rule50_count: " << static_cast(entry.rule50_count) + << "\n"; + std::cout << " invariance_info: " << static_cast(entry.invariance_info) + << " (" << DecodeInvarianceInfo(entry.invariance_info) << ")\n"; + std::cout << " dummy: " << static_cast(entry.dummy) << "\n"; + std::cout << " root_q: " << entry.root_q << "\n"; + std::cout << " best_q: " << entry.best_q << "\n"; + std::cout << " root_d: " << entry.root_d << "\n"; + std::cout << " best_d: " << entry.best_d << "\n"; + std::cout << " root_m: " << entry.root_m << "\n"; + std::cout << " best_m: " << entry.best_m << "\n"; + std::cout << " plies_left: " << entry.plies_left << "\n"; + std::cout << " result_q: " << entry.result_q << "\n"; + std::cout << " result_d: " << entry.result_d << "\n"; + std::cout << " played_q: " << entry.played_q << "\n"; + std::cout << " played_d: " << entry.played_d << "\n"; + std::cout << " played_m: " << entry.played_m << "\n"; + std::cout << " orig_q: " << entry.orig_q << "\n"; + std::cout << " orig_d: " << entry.orig_d << "\n"; + std::cout << " orig_m: " << entry.orig_m << "\n"; + std::cout << " visits: " << entry.visits << "\n"; + std::cout << " played_idx: " << entry.played_idx << "\n"; + std::cout << " best_idx: " << entry.best_idx << "\n"; + std::cout << " policy_kld: " << entry.policy_kld << "\n"; + PrintFloatArray(entry.probabilities, std::size(entry.probabilities), + "probabilities", float_per_line); + PrintUint64Array(entry.planes, std::size(entry.planes), "planes", + plane_per_line); + std::cout << std::flush; +} + +} // namespace training +} // namespace lczero diff --git a/csrc/utils/training_data_printer.h b/csrc/utils/training_data_printer.h new file mode 100644 index 00000000..174d5e28 --- /dev/null +++ b/csrc/utils/training_data_printer.h @@ -0,0 +1,38 @@ +#ifndef LCZERO_TRAINING_UTILS_TRAINING_DATA_PRINTER_H_ +#define LCZERO_TRAINING_UTILS_TRAINING_DATA_PRINTER_H_ + +#include + +#include +#include +#include + +#include "loader/frame_type.h" +#include "trainingdata/trainingdata_v6.h" + +namespace lczero { +namespace training { + +// Prints a float array with configurable number of values per line. +void PrintFloatArray(const float* data, size_t size, absl::string_view name, + int64_t per_line); + +// Prints a uint64 array with configurable number of values per line. +void PrintUint64Array(const uint64_t* data, size_t size, absl::string_view name, + int64_t per_line); + +// Decodes the invariance_info byte into a human-readable string. +std::string DecodeInvarianceInfo(uint8_t invariance_info); + +// Converts a V6TrainingData entry to FEN (Forsyth-Edwards Notation). +std::string TrainingDataToFen(const V6TrainingData& entry); + +// Prints a V6TrainingData entry with a custom header and formatting options. +void PrintTrainingDataEntry(const FrameType& entry, + absl::string_view header_text, + int64_t float_per_line, int64_t plane_per_line); + +} // namespace training +} // namespace lczero + +#endif // LCZERO_TRAINING_UTILS_TRAINING_DATA_PRINTER_H_ diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 00000000..4ad47b70 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,287 @@ +# Running "new" training pipeline + +Note that the code is still in active development, so things change a lot. +The current document was last updated on 2025-11-30. + +## Building + +The new training pipeline is located in `src/` (Python part) and `csrc/` (C++ +part). + +* Python code uses `uv`. Install it as described in the + [uv installation guide](https://docs.astral.sh/uv/#installation). +* Many steps are run via `just`. `just` is just a glorified shell script runner. + So either look into [`Justfile`](../justfile) or install `just` as described + in the [just installation guide](https://github.com/casey/just#installation). +* You'll need a recent protobuf compiler (`protoc`). +* You'll need a C++ compiler. In this example we use `clang`. + +```bash +cd +uv python install 3.12 +uv venv +uv sync +uv pip install meson ruff +git submodule update --init --recursive +CXX=clang++ CC=clang uv run meson setup build/release\ + --buildtype=release --native-file=native.ini +uv run meson configure build/release \ + -Dcpp_args='-Wno-error=deprecated-declarations' +just build +cd src/lczero_training +ln -sfT ../../build/release/_lczero_training.cpython-*-x86_64-linux-gnu.so _lczero_training.so +just build-proto +``` + +## Training a model + +To train a model you need: + +* Training data +* A configuration file +* Create a checkpoint. +* Run the pipeline. + +### Training data + +Unlike the old training pipeline, the new one doesn't need .tar files to be +unpacked. While it does support plain `.gz` chunk files, it's not efficient as +it stores each individual file name in memory. So instead, use `.tar` files, +the tool can index and seek inside them. + +The tool watches a directory (and its subdirectories) for new files. + +Terms used: + +* **Chunk**/Game: A single training game, individual `.gz` file. +* **Chunk source**: A file (`.tar` or `.gz`) containing multiple chunks. +* **Frame**/Record/Position: A single training position inside a chunk. +* **Training tensor**: A single batch of inputs/outputs encoded in NN format for + one training step. + +Incoming data comes as chunks, but for the training we need frames from +different games. + +### Note on RL vs SL training + +The tool supports both supervised learning (SL) and reinforcement learning (RL). +Here is overview of the configuration differences: + +| RL Training | SL Training | +| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `shuffling_chunk_pool` has relatively small`chunk_pool_size`, which would be used as a sliding window. | `chunk_pool_size` should be larger than all data, so that all data is used for training. | +| `training.schedule.chunks_per_network` is non zero — that's how many new chunks to wait for before starting a new epoch. | `chunks_per_network` is zero. Once an epoch is done, it starts a new one immediately. | +| RL currently uses "hanse sampling", where currently the entire chunk is loaded and rescored to use just one position from it. The reservoir `shuffling_frame_sampler` is not used in this case. This is currently slow (until we implement caching), so it limits the throughput. | For SL, it better to use two stage sampling: `shuffling_chunk_pool` in non-hanse mode, and then `shuffling_frame_sampler` to shuffle positions within chunks. | + +### Creating a checkpoint + +To create a fresh checkpoint, you'll need `model` and `training` sections of the +configuration file to be filled. + +Then run: + +```bash +uv run lc0-init --config .textproto --lczero_model .pb.gz +``` + +The `--lczero_model` parameter is optional. If not given, the network is +initialized with random weights. + +### Training + +Run: + +```bash +CUDA_VISIBLE_DEVICES=0 uv run lc0-tui --config .textproto --logfile train.log +``` + +Notes: + +* There are log files both in `tui` and in the configuration file. They +are slightly different, TUI log usually is more useful. +* In the tool, you can press `q` to quit, and `Ctrl+p` to get a command palette. +There, you have one useful command: "Start training immediately". +* For multi-GPU training, ensure your batch size is divisible by the number of GPUs. +* The overfit utility (`uv run overfit`) does not support multi-GPU. Set `CUDA_VISIBLE_DEVICES` to use only one GPU when running overfit. +* Also note that TUI is 100% vibe coded, so you'll see lots of mocks in the UI. +:-P + +## Tools + +The repository consists of set of tools, mostly written in Python, but some are +in C++. To run a Python tool, use `uv run `. C++ tools are binaries in +`build/release`. Most tools need a configuration file as a parameter (see +below). + +### Python tools + +| Tool | Description | +| ---------------------- | ------------------------------------------------------------------------------------------------------------ | +| lc0-daemon | The main training daemon. It acts as a JSONL server, so it's not usable directly from command line yet. | +| lc0-tui | A terminal user interface that runs the training daemon. Here is what you have to run. | +| lc0-init | Initializes a new training run/checkpoint. | +| lc0-migrate-checkpoint | Migrates JAX/Orbax checkpoint after model/training configuration changes. | +| lc0-overfit | Runs an overfitting test: takes one batch from the data loader and repeatedly trains on it | +| lc0-eval | Evals batches from the data loader on a given checkpoint, can dump inputs/outputs in various formats. | +| lc0-leela2jax | | +| lc0-describe | | +| lc0-test-dataloader | | +| lc0-tune-lr | Trains on exponentially increasing learning rate, and outputs losses into csv file. Useful for picking a LR. | +| lc0-backfill-metrics | Loads older checkpoints computes metrics for them, and exports them to tensorboard. | +| lc0-train | Trains a single epoch (doesn't save or export the model though). Used for benchmarking. | +| lc0-weights | Manipulates weight files: arithmetic operations, grafting components, format conversion. See [weights_tool.md](weights_tool.md). | + +### C++ tools + +| Tool | Description | +| ---------------------------- | --------------------------------------------------------- | +| rescore_chunk | Runs rescorer on a single chunk | +| startpos_policy_distribution | | +| result_distribution | | +| filter_chunks | | +| dump_chunk | Dumps the content of a chunk file for debugging purposes. | + +## Configuration + +The configuration is a text protobuf file, with the following sections: + +| Section | Description | +| ----------- | ---------------------------------------------- | +| data_loader | Configuration for the data loader. | +| model | Model architecture configuration. | +| training | Configuration for the training configuration. | +| metrics | Metrics to export into tensorboard. | +| export | Configuration for exporting the trained model. | + +Also it has `log_filename` field (where to write the log) and `name` field (must +be `little-teapot`). + +It's recommended to use existing configuration files as a starting point. + +### Data loader configuration + +Data loader is a pipeline that consists of pluggable stages. Here are stages +that are currently implemented: + +| Stage type | Description | Input | Output | +| ------------------------- | ------------------------------------------------------------------------------------------------ | ------------ | --------------------- | +| `file_path_provider` | Watches a directory for existing and new files. | None | Filenames | +| `chunk_source_reader` | Reads and indexes chunk source files (`.tar` or `.gz`). | Filenames | ChunkSources | +| `chunk_source_splitter` | Splits chunk sources into smaller chunk sources give the proportion (used for test/train split). | ChunkSources | multiple ChunkSources | +| `shuffling_chunk_pool` | Accumulates chunk sources and outputs chunks in shuffled order | ChunkSources | Chunks | +| `simple_chunk_extractor` | Unpacks chunks from chunk sources | ChunkSources | Chunks | +| `chunk_rescorer` | Rescores chunks | Chunks | Chunks | +| `chunk_unpacker` | Extracts positions from chunks | Chunks | Frames | +| `shuffling_frame_sampler` | Outputs frames in shuffled order | Frames | Frames | +| `tensor_generator` | Converts frames into training batches in numpy tensor format | Frames | Training Tensors | + +The pipeline ends with one or more outputs, which provide tuples of batched +tensors for training. + +> [!NOTE] +> The current format of the training batch is +> +> * `inputs`: float32 tensor of shape `[batch_size, 112, 8, 8]` +> * `policy_target`: float32 tensor of shape `[batch_size, 1862]` +> * `value_target`: float32 tensor of shape `[batch_size, 6, 3]`, where 6 rows +> are sources of the value (`result`, `best`, `played`, `orig`, `root` and +> `st`), and 3 columns are (`q` (w-l), `draw`, `movesleft`). + +Every stage must have an unique name (may or not be the same as the stage type), +and arbitrary number of inputs (depending on the stage type; most have one +input). + +Here is the structure of the data loader configuration: + +```textproto +stage { + name: "file_provider" + file_path_provider { + # ... + } +} +stage { + name: "loader" + input: "file_provider" + chunk_source_reader { + # ... + output { name: "myoutput" } + } +} +stage { + name: "chunk_shuffler" + input: "loader.myoutput" + shuffling_chunk_pool { + # ... + } +} +# ... +stage { + name: "tensor_gen" + input: "sampler" + tensor_generator { + batch_size: 256 + # ... + } +} +output: "tensor_gen" # unnamed output +output: "test:test_tensor_gen" # named output +``` + +#### Stage output configuration + +Every stage provides one or more outputs. The configuration of the output is like this (all fields optional): + +```textproto +output { + name: "myoutput" + queue_capacity: 8 # default: 4 + overflow_behavior: BLOCK +} +``` + +* By default, outputs are not named, but you can name them. +* Higher `queue_capacity` allows you to "pre-cache" data, so that when the rate + of the producer stage is spiky, the pipeline is not blocked. On the other + hand, the data in the queue may be "stale" (i.e. when you train a new network, + the data in the queue is still for the old network). +* `overflow_behavior` controls what happens when the output queue is full: + * `BLOCK`: default and what's needed for most stages. The producer stage is + blocked until there is space in the queue. + * `DROP_NEW` and `KEEP_NEWEST` drops the data from the queue (either the + incoming data, or the oldest data in the queue). These are useful e.g. for + auxiliary output of a stage (e.g. validation), so that the auxiliary + pipeline doesn't block the main pipeline. + +### Stage configurations + +#### file_path_provider + +Watches a directory for existing and new files. First it sends all existing +files, then sends special "Initial Scan Done" event, and then watches for new +files. + +#### chunk_source_loader + +Takes the filenames from the input, and loads them as chunk sources. Skips files +which are not chunk sources. + +* `frame_format`: `V6TrainingData` (default) or `V7TrainingData`. + +#### shuffling_chunk_pool + +Shuffling chunk pool is the central part of the data loader. In most cases, +it is the only stage responsible for shuffling the data. In some cases, you may +want to have secondary shuffling_frame_sampler after it (e.g. for SL training). + +Every chunk source has a "sort key" (currently, it's the file name without +path). It's needed to determine the order of chunks to use for the sliding +window. + +* `chunk_pool_size`: The size of the training window, in number of chunks. Even + when there are not enough chunks yet, the stage will output chunks from what + it has. It will not start producing data until the "initial scan done" event + is received from the file_path_provider. + * For RL training, typical values are 250k to 5M. + * For SL training, it should be larger than all data, so that all data is used + for training. \ No newline at end of file diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 00000000..aadb3aa7 --- /dev/null +++ b/docs/architecture.md @@ -0,0 +1,78 @@ +# Architecture Overview + +The document outlines the architecture of the new Leela Chess Zero training +system. The training process involves a Reinforcement Learning (RL) pipeline +where new training data is continuously generated. + +The old script required fresh starts to train a new network on fresher data. +Furthermore, the use of TensorFlow involved a very long model compilation time +(approx. 1 hour), which dominated the actual training time for a single epoch. + +The new script will be a single, long-running Python application that automates +the entire cycle: + +1. Monitors for a sufficient amount of new training data. +2. Triggers and executes the training of a new network. +3. Exports the trained network for use. + +The core training loop will be implemented in JAX. The data +[loading and preprocessing pipeline](loader.md) is a C++ code exposed to Python +via pybind11. This C++ library is internally multi-threaded but exposes a simple +API to Python (GetNextBatch(), GetStats()). + +We'd like to have a fancy TUI dashboard to monitor the loading and training +process. + +## TrainingDaemon + +The main (in terms of importance) class of the training code is +`TrainingDaemon`, which: + +* Is a jsonl server operating through stdin/stdout, implementing the protocol + described in [JSONL IPC Protocol](jsonl.md). + * Receives a command to start training with the location of the config file. + * Other commands are possible in the future. + * Sends periodic progress notifications. +* Owns the data loader. +* Waits for the data loader to ingest enough new chunks. +* Starts the training loop when enough data is available. +* Finalizes and uploads the trained network. + +## Configuration + +The configuration is a large nested dataclass structure, which covers: + +* [Data loader](../src/lczero_training/config/data_loader_config.py) to be + passed to the C++ data loader. +* Information for the training daemon, e.g. how many chunks to wait for + before starting the training. +* Model definition, for model builder. +* Training parameters, such as batch size, number of epochs, etc. +* Export parameters, such as the path to export the trained model to. + +From user perspective, the configuration is a YAML file, which is parsed +into the dataclass structure. + +## Data Loader + +The internals of the data loader are described in detail in [Data Loader](loader.md). + +From python perspective, it has the following interface: + +* Constructor takes a `DataLoaderConfig` dataclass. +* `GetNextBatch()` returns a tuple of buffer-protocol-compliant tensors. + * Later it will have a parameter that specifies wether we need training, test + or validation batch. +* `GetStats()` returns a DataClass (exact structure TBD) with the current + statistics of the data loader. + +## TUI + +TUI is a separate frontend app, implemented as `TrainingTuiApp`, which runs +`TrainingDaemon` as a subprocess (and communicates through jsonl via +stdin/stdout). + +* TUI is `Textual`-based app. +* Located in `src/lczero_training/tui/`. +* UI ideas are described in [TUI](tui.md). +* TUI will have a log pane which shows stderr of `TrainingDaemon`. diff --git a/docs/checkpoint_migration.md b/docs/checkpoint_migration.md new file mode 100644 index 00000000..3b458409 --- /dev/null +++ b/docs/checkpoint_migration.md @@ -0,0 +1,129 @@ +# Checkpoint Migration + +When part of the model or training setup changes, JAX training state checkpoints +may become incompatible with the new setup. + +For this, we provide a utility to help migrate checkpoints to the new setup. + +The underlying implementation lives in +`src/lczero_training/training/migrate_checkpoint.py` and is exposed via the CLI +command `migrate-checkpoint` registered in `pyproject.toml`. + +## Command line arguments + +* `--config`: Path to the RootConfig textproto config. `model`/`training` + sections of this config will be used to initialize the new training state. + `checkpoint` is the location of the old checkpoint to migrate. +* `--new_checkpoint`: Path to save the new checkpoint to. If not set, the tool + only checks whether the migration rules fully cover the differences between + the old and new training states. +* `--overwrite`: If set, allows overwriting existing checkpoint. Also in this + case, if `--new_checkpoint` is not set, old checkpoint is used as the new + checkpoint. +* `--rules_file`: Path to a CheckpointMigrationConfig textproto file containing + the migration rules. See below for the format of this file. If not set, no + migration rules will be applied (used for debugging, or to check that the + old and new states are identical). +* `--serialized-model`: If set, use serialized state for a model. Checkpoint + already loads serialized as we do not provide schema. This is needed to avoid + `GetAttrKey`s. +* `--checkpoint_step`: If set, use this step when loading from old checkpoint + instead of the latest. +* `--new_checkpoint_step`: If set, use this step when saving the new checkpoint + instead of copying the old step. + +## Creation of the state to compare + +The checkpoint is loaded into a raw pytree, i.e. template is not passed to it. +(old checkpoint with new model would fail to load). + +Roughly, like this: + +```python + manager = ocp.CheckpointManager( + filepath, + options=ocp.CheckpointManagerOptions(create=False), + ) + state = manager.restore(step=None) +``` + +The model state is created using `TrainingState.new_from_config` +(`src/lczero_training/training/state.py`). If `--serialized-model` is passed, +the model state is serialized using `flax.serialization.to_state_dict`. + +## Migration rules format + +The migration rules are specified in a `CheckpointMigrationConfig` textproto +file. `CheckpointMigrationConfig` is defined in +`proto/checkpoint_migration_config.proto` + +It contains a list of `CheckpointMigrationRule` messages. + +Every rules has a `from_path` and `to_path` field (both optional, but at least +one must be set). These are string fields, which are json list of strings and +integers, and are mapped to pytree KeyPaths: + +* Integers are used as `SequenceKey` (list/tuple indices). +* Strings are used as `DictKey` (dict keys). +* If other types is met in the path, an error is raised. +* If other key types are met in `PathKey`, the error is also raised. + +* By default, all keys that are present both in the old and new state are + preserved (copied from old to new). +* If both `from_path` and `to_path` are set, they must be different. The old + values at `from_path` are copied to `to_path` in the new state. +* If only `to_path` is set, it means that the new state at `to_path` is + taken from the new initialized state. +* If only `from_path` is set, it means that the old state at `from_path` is not + present in the new state, and is ignored. Without this rule, the migration + would fail if the old state has keys that are not present in the new state. +* If we want to keep a initialized ("new") value at a path that is also + present in the old state, we use two rules: one with only `to_path` to + initialize it, and one with only `from_path` to ignore the old value at + that path. + +Note that the `from_path` and `to_path` are prefixes of the actual paths, so the +actual subtrees are copied/ignored. + +## Working of the migration + +1. Both new and old states are created. +2. Both of them are flattened using `jax.tree_util.tree_flatten_with_path` into + list of `(KeyPath, value)` pairs and a treedef. +3. Lists of `(KeyPath, value)` are converted to dicts mapping `KeyPath` to + `value`. +4. We build a set of `source_paths` (keys of the old state). +5. We build a set of `dest_paths` (keys of the new state). +6. Rules are applied in arbitrary order as following: + * If both `from_path` and `to_path` are set, the values prefixed by + `from_path` are copied to corresponding `to_path` in the new state. Of + course, `to_path` in the new state must exist (i.e. we never create new + keys). The copied source paths are removed from `source_paths`. + The corresponding destination paths are removed from `dest_paths`. + * If only `to_path` is set, we delete all paths prefixed by `to_path` from + `dest_paths`. This means that we keep the initialized value at this path. + * If only `from_path` is set, we delete all paths prefixed by `from_path` + from `source_paths`. This means that we ignore the old value at this path. +7. After all rules are applied, we check that `source_paths` and `dest_paths` + are equal. If not, the migration is incomplete, and we raise an error. +8. After this, we copy all remaining `source_paths` (i.e. those that were not + mentioned in any rule) to the new state. This means that by default, all + keys that are present both in the old and new state are preserved (copied + from old to new). +9. Finally, we unflatten the new state dict back to a pytree using the new + treedef. + +If `--new_checkpoint` is set, we save the new state to the specified path. +Otherwise, we just print that the migration is possible with the given rules. + +Instead of just printing one error, the tool should print ALL errors it finds. +If should be helpful to fix the rules. + +## Implementation notes + +* There is `justfile`, e.g. `just build-proto` to build the protos. +* We use `uv`. Example usage: + +```bash +uv run migrate-checkpoint --config=~/tmp/lc0/config/overfit.textproto +``` diff --git a/docs/example.textproto b/docs/example.textproto new file mode 100644 index 00000000..90ad5199 --- /dev/null +++ b/docs/example.textproto @@ -0,0 +1,189 @@ +# Example configuration file for lczero-training +# This file demonstrates all available configuration options with their default +# values and explanations of what each setting controls. + +name: "little-teapot" +data_loader { + stage { + name: "file_path_provider" + file_path_provider { + # Directory with training data files. + directory: "/home/crem/tmp/2025-07/lczero-training/data2" + output { + queue_capacity: 16 # Internal file queue size + } + } + } + stage { + name: "chunk_source_loader" + input: "file_path_provider" + chunk_source_loader { + threads: 1 # Threads for loading chunks + frame_format: V6TrainingData # Training data format (V6TrainingData or V7TrainingData) + output { + queue_capacity: 16 # Output queue for chunk sources + } + } + } + stage { + name: "shuffling_chunk_pool" + input: "chunk_source_loader" + shuffling_chunk_pool { + chunk_pool_size: 50000 # Shuffle buffer size (chunks in memory) + source_ingestion_threads: 1 # Threads for ingesting new sources + chunk_loading_threads: 4 # Threads for loading chunk data + output { + queue_capacity: 16 # Output queue for shuffled chunks + } + } + } + stage { + name: "chunk_rescorer" + input: "shuffling_chunk_pool" + chunk_rescorer { + threads: 1 # Threads for chunk rescoring + output { + queue_capacity: 16 # Output queue for rescored chunks + } + syzygy_paths: "/path/to/tb" # Tablebase search paths (comma-separated) + dist_temp: 1.0 # Policy temperature applied during rescoring + dist_offset: 0.0 # Policy offset applied during rescoring + dtz_boost: 0.0 # DTZ boost for endgame policy tuning + new_input_format: -1 # Keep original input format (-1 disables change) + deblunder_threshold: 0.10 # Threshold for policy deblundering adjustments + deblunder_width: 0.06 # Width controlling smoothing around threshold + } + } + stage { + name: "chunk_unpacker" + input: "chunk_rescorer" + chunk_unpacker { + threads: 1 # Threads for unpacking chunks + # Probability of sampling each position within a chunk. + position_sampling_rate: 0.03 + output { + queue_capacity: 16 # Output queue for unpacked frames + } + } + } + stage { + name: "shuffling_frame_sampler" + input: "chunk_unpacker" + shuffling_frame_sampler { + threads: 1 # Threads for frame sampling + reservoir_size_per_thread: 1000000 # Sampling reservoir per thread + output { + queue_capacity: 16 # Output queue for sampled frames + } + } + } + stage { + name: "tensor_generator" + input: "shuffling_frame_sampler" + tensor_generator { + threads: 1 # Threads for tensor generation + batch_size: 128 # Batch size for tensors + output { + queue_capacity: 8 # Output queue for batched tensors + } + } + } + output: "tensor_generator" +} +model { + defaults { + compute_dtype: F32 + activation: ACTIVATION_MISH + ffn_activation: ACTIVATION_MISH + } + embedding { dense_size: 512 embedding_size: 1024 dff: 1536 } + encoder { + num_blocks: 15 + dff: 1536 + d_model: 1024 + heads: 32 + smolgen { + hidden_channels: 32 + hidden_size: 256 + gen_size: 256 + activation: ACTIVATION_SWISH + } + } + policy_head { name: "vanilla" embedding_size: 1024 d_model: 1024 } + value_head { name: "winner" num_channels: 128 } + movesleft_head { name: "main" num_channels: 32 } +} +training { + schedule { + steps_per_network: 250 + chunks_per_network: 50000 + } + lr_schedule { + starting_step: 0 + duration_steps: 0 # 0 means indefinite duration + lr: 0.001 + } + # Example multi-phase schedule with warmup: + # lr_schedule { + # starting_step: 0 + # duration_steps: [1500, 500, 0] # Warmup, transition, then indefinite + # lr: [0.0, 0.0005, 0.0005] # Start at 0, ramp to 0.0005, hold + # transition: [LINEAR, CONSTANT] # Linear warmup, then constant + # # Missing transitions default to CONSTANT. Last duration 0 = indefinite. + # } + checkpoint { + path: "/home/crem/tmp/2025-09/lc0_training/checkpoint" + max_to_keep: 5 + } + optimizer { + nadamw { + beta_1: 0.9 beta_2: 0.98 epsilon: 1e-7 weight_decay: 0.0001 + # Rule order matters: first match wins. + decay_selector { + rule { match: "**/bias" include: false } + rule { match: "**/ln*/**" include: false } + rule { match: "**/embedding/embedding/**" include: false } + rule { match: "**/policy_heads/**" include: true } + rule { match: "**/value_heads/**" include: true } + rule { match: "**/movesleft_heads/**" include: true } + otherwise_include: false + } + } + # Alternative optimizers: + # nadam { beta_1: 0.9 beta_2: 0.999 epsilon: 1e-8 } + # sgd { momentum: 0.9 nesterov: true } + # freeze_selector freezes matching weights (they receive no gradient + # updates). + # freeze_selector { + # rule { match: "**/embedding/**" include: true } + # otherwise_include: false + # } + } + max_grad_norm: 10.0 # Global gradient-norm clip; omit or set to 0 to disable. + losses { + policy { + head_name: "vanilla" + metric_name: "main_ce" + weight: 1.0 + illegal_moves: MASK + type: CROSS_ENTROPY + } + policy { + head_name: "vanilla" + metric_name: "main_kl" + weight: 1.0 + illegal_moves: MASK + type: KL + temperature: 1.0 # Softmax temperature applied before KL evaluation + } + value { head_name: "winner" weight: 1.0 } + movesleft { head_name: "main" weight: 1.0 } + } +} +metrics { + tensorboard_path: "/tmp/tensorboard/myrun" +} +export { + destination_filename: "/home/crem/tmp/2025-08/lc0_training/exported_models/lc0-{datetime}-{step:08d}.pb.gz" + upload_training_run: 3 +} diff --git a/docs/heads.md b/docs/heads.md new file mode 100644 index 00000000..1f1bfc0a --- /dev/null +++ b/docs/heads.md @@ -0,0 +1,100 @@ +# Neural Network Heads Documentation + +This document describes the various policy and value heads used in the network, their training targets, and any specific scaling or transformations applied during training. + +## Policy Heads + +### 1. Vanilla Policy (`vanilla`) +* **Description**: The standard policy head predicting the best move. +* **Training Target**: The `probabilities` vector from the training data (MCTS visit counts). +* **Scaling/Transformation**: No specific scaling is applied to the target. The loss function compares the network output (logits) directly against the target probability distribution. +* **Loss Function**: Cross-entropy loss (Kullback-Leibler divergence). + +### 2. Optimistic Short-Term Policy (`optimistic_st`) +* **Description**: A policy head trained to be "optimistic" about the outcome, focusing more on moves that lead to better short-term evaluations. It uses a weighted loss function where positions that the network underestimates (target > prediction) are weighted more heavily. +* **Training Target**: The same `probabilities` vector as the `vanilla` head. +* **Scaling/Transformation**: + * **Mechanism**: The "optimism" is applied as a **sample weight** in the loss function. + * **Computation Guide**: + 1. **Inputs**: + * `v_st_target`: Short-term value target (scalar, from `st` head target). + * `v_st_pred`: Short-term value prediction (scalar, from `st` head output). + * `v_st_err_pred`: Predicted squared error of the short-term value (scalar, from `st_err` head output). + 2. **Standard Deviation Estimation**: + * `sigma = sqrt(v_st_err_pred)` + 3. **Z-Score Calculation**: + * `z = (v_st_target - v_st_pred) / (sigma + 1e-5)` + * This measures how many standard deviations the target is away from the prediction. A positive `z` means the position is better than predicted (underestimated). + 4. **Weight Calculation**: + * `strength = 2.0` (default configuration) + * `weight = sigmoid((z - strength) * 3)` + * **Interpretation**: + * If `z` is large (positive), meaning the target is much higher than predicted (highly underestimated), the weight approaches 1. + * If `z` is small or negative (overestimated or accurately predicted), the weight approaches 0. + * The `strength` parameter shifts the sigmoid, controlling the threshold of "optimism" required to trigger training. +* **Loss Function**: Weighted Cross-entropy loss. `loss = weight * CrossEntropy(target, output)`. + +### 3. Soft Policy (`soft`) +* **Description**: A policy head trained on a "softened" version of the MCTS probabilities, encouraging exploration or capturing more of the distribution's shape. +* **Training Target**: The `probabilities` vector from the training data. +* **Scaling/Transformation**: + * **Mechanism**: **Temperature scaling** is applied to the target probabilities **inside the loss function**. + * **Calculation**: `target = target^(1/temperature)`. + * **Location**: This transformation happens in the loss calculation logic (specifically in `correct_policy` helper in `tfprocess.py`), **not** at the tensor generation stage. The tensor generator provides the raw `probabilities`. +* **Loss Function**: Cross-entropy loss against the temperature-scaled target. + +### 4. Opponent Policy (`opponent`) +* **Description**: Predicts the move the opponent actually played in the game. +* **Training Target**: `opp_played_idx` (Integer index of the move played by the opponent). +* **Scaling/Transformation**: The integer index is converted to a one-hot vector inside the loss function. +* **Loss Function**: Cross-entropy loss. + +## Value Heads + +### 1. Winner (`winner`) +* **Description**: Predicts the final game outcome (Win/Draw/Loss). +* **Training Target**: A 3-element probability vector derived from `result_q` and `result_d` in the training data. + * `Win = (1 + result_q - result_d) / 2` + * `Loss = (1 - result_q - result_d) / 2` + * `Draw = result_d` +* **Scaling/Transformation**: None. +* **Loss Function**: Cross-entropy or MSE depending on configuration. + +### 2. Q-Value (`q`) +* **Description**: Predicts the expected value of the position based on the MCTS search (best Q). +* **Training Target**: A 3-element probability vector derived from `best_q` and `best_d` in the training data. + * `Win = (1 + best_q - best_d) / 2` + * `Loss = (1 - best_q - best_d) / 2` + * `Draw = best_d` +* **Scaling/Transformation**: None. +* **Loss Function**: MSE or Cross-entropy. + +### 4. Q-Value Error (`q_err`) +* **Description**: Predicts the squared error of the `q` head prediction compared to the target. +* **Training Target**: `(q_target - q_pred)^2`. +* **Scaling/Transformation**: None. +* **Loss Function**: MSE. + +### 5. Short-Term Value (`st`) +* **Description**: Predicts the short-term evaluation of the position (e.g., from a shallow search or static eval). +* **Training Target**: A 3-element probability vector derived from `q_st` and `d_st` in the training data. + * `Win = (1 + q_st - d_st) / 2` + * `Loss = (1 - q_st - d_st) / 2` + * `Draw = d_st` +* **Scaling/Transformation**: None. +* **Loss Function**: MSE or Cross-entropy. + +### 6. Short-Term Value Error (`st_err`) +* **Description**: Predicts the squared error of the `st` head prediction compared to the target. Used for calculating uncertainty/variance for the `optimistic_st` head. +* **Training Target**: `(st_target - st_pred)^2`. +* **Scaling/Transformation**: None. +* **Loss Function**: MSE. + +## Auxiliary Heads + +### 1. Moves Left (`moves_left`) +* **Description**: Predicts the number of moves remaining in the game. +* **Training Target**: `plies_left` (Float). +* **Scaling/Transformation**: + * **Mechanism**: The target and output are scaled down by a factor (e.g., 20.0) **inside the loss function** to bring the loss magnitude into a similar range as other losses. +* **Loss Function**: Huber loss. diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 00000000..12f5679a --- /dev/null +++ b/docs/index.md @@ -0,0 +1,10 @@ +# Index + +* [Overview, glossary and file formats](overview.md) — A an overview of the + project, including definitions of some terms and file formats used. +* [Data Loader](loader.md) — A module for loading, preprocessing, shuffling, and + feeding training data. +* [Training Tuple Format](training_tuple.md) — A description of the training + tuple format used in the project. This is an interface between data loader and + model training. +* [Command-Line Tools](cli.md) — How to run the tools via `uv run`. diff --git a/docs/loader.md b/docs/loader.md new file mode 100644 index 00000000..9bdd20fc --- /dev/null +++ b/docs/loader.md @@ -0,0 +1,220 @@ +# Data Loader + +The Data Loader is a C++ module (exposed to Python via pybind11) that handles +loading, preprocessing, shuffling, and feeding training data for the Leela Chess +Zero training process. + +## Python Integration + +The loader has been exposed to Python through pybind11, allowing direct use of +the C++ `DataLoader` from Python code. Key aspects: + +* **Configuration**: Generated protobufs (for example `DataLoaderConfig`) are + passed directly to the binding, or via the convenience wrapper + `lczero_training.dataloader.make_dataloader`. +* **Control Plane**: Use `DataLoader.send_control_message()` with + `proto.stage_control_pb2.StageControlRequest` to fan out commands such as + chunk-pool anchor updates. +* **Memory Management**: Uses `unique_ptr::release()` with `py::return_value_policy::take_ownership` for efficient tensor ownership transfer +* **Output Format**: Returns tuple of numpy arrays compatible with JAX through the buffer protocol +* **Usage**: `from lczero_training.dataloader import make_dataloader` + +## High-Level Overview + +The Data Loader consists of the following stages connected through a +[Queue](../csrc/utils/queue.h): + +* [FilePathProvider](../csrc/loader/chunk_feed/file_path_provider.h) — Training + data discovery worker (watches a directory and provides feed of filenames) +* [ChunkSourceLoader](../csrc/loader/chunk_feed/chunk_source_loader.h) — Reads + chunks from files, providing a stream of chunks. +* [ShufflingChunkPool](../csrc/loader/chunk_feed/shuffling_chunk_pool.h) — Keeps + a set of chunks, managing the last `num_chunks` available and removing old + ones, and outputting them in shuffled order. +* (skip for now) [ChunkValidator](../csrc/loader/chunk_feed/chunk_validator.h) — + Filters the chunk stream, filtering out invalid chunks. +* [ChunkRescorer](../csrc/loader/stages/chunk_rescorer.h) — Rescores chunks + using Syzygy tablebases and configurable policy adjustments. +* [ChunkUnpacker](../csrc/loader/chunk_feed/chunk_unpacker.h) — Unpacks + chunks into frames, which are then processed by the next stages. +* [ShufflingFrameSampler](../csrc/loader/shuffling_frame_sampler.h) — Takes a + stream of frames and provides shuffled batches of frames for training, using + reservoir sampling. +* [TensorGenerator](../csrc/loader/tensor_generator.h) — Takes frames and + provides tensor buffers for the training process. + +## Metrics + +All stages expose the following metrics in +[DataLoaderMetricsProto](../proto/training_config.proto): + +* load — for measure how much time the threads are working vs idle. +* queue - for monitoring queue statistics. + +There are the following exceptions: + +* ShufflingChunkPool + * Has two thread pools (indexing and chunk loading), so needs two `load` + metrics. + * Needs metric (statisticsmetric) for current number of chunk sources. + * Needs metric (simple value) for current number of chunks in the pool. + * Needs metric (simple value) for pool capacity. + +* ChunkUnpacker + * Needs to track the number of bad chunks (statisticsmetric) + +* ShufflingFrameSampler + * Needs capacity of reservoir (simple value) + * Needs current size of reservoir (simple value) + +### Adding a new metric + +To add a new metric, you need to: + +* Add a field in the relevant proto in + [training_config.proto](../proto/training_config.proto) +* Add a field in a relevant stage to collect the metric (if needed). +* Implement FlushMetrics() function in your stage. It should reset internal + state metrics to zero, and return what it accumulated (or latest value). +* In [DataLoader::MetricsThread](../csrc/loader/data_loader.cc) call + FlushMetrics() of the relevant stage and assign it to proper proto field. +* In the proper + [Queue or Stage widget](../src/lczero_training/tui/stage_widgets.py) create + proper UI elements and update update_metrics() function to update them. + +* For load metrics, you have to create LoadMetricPauser per thread, e.g. see [ShufflingChunkPool](../csrc/loader/chunk_feed/shuffling_chunk_pool.cc). +* For queue metrics, just collect the queue statistics in FlushMetrics(). +* For all metrics, both "all time" and "during last second" is provided. The + latter is useful to determine rate (items per second). +* In general, use StatisticsProtoInt64 (or StatisticsProtoDouble) for + distribution metrics, and simple number field for additive metrics like + counts. + +## TensorGenerator + +Batch size is configurable in the stage options. + +The `TensorGenerator` stage takes frames from the input queue and produces tuple +of tensors for tensor returned as [TensorTuple](../csrc/utils/tensor.h). +The first dimension of every tensor in the tuple is the batch size, and the +rest are described in the [Training Tuple Format](training_tuple.md). + +## Stage interface + +All stages implement the similar API and structure, although not sharing any +base class. + +All stages/workers (except pure producers) wait for the input queue to Close(), +then Close() output Queue. + +```cpp +class Stage { + public: + using InputType = ...; // Type of input data for this stage + using OutputType = ...; // Type of output data from this stage + // input_queue is omitted in the producer stages like FilePathProvider. + Stage(Queue* input_queue, /* other params */); + Queue* output(); + +private: + ThreadPool thread_pool_; + Queue* input_queue_; + Queue output_queue_; +}; +``` + +## Chunk Set + +The Chunk Set takes the feed of chunk sources, indexes them, and assigns the +chunk range (base; base+num_chunks) to each chunk source. It aims to keep the +newest chunk sources that cover the last `chunks_window_` chunks, and removes +old chunk sources when new ones are added. + +On the output side, it returns a stream of chunks within +(`last - chunk_window_`, `last`) range, without repetitions. To do that, it +utilizes a [StreamShuffler](../csrc/loader/stream_shuffler.h) to which provides +the shuffled stream of numbers within the (dynamic) range. + +The Chunk Set gets the stream of chunks (initial chunks are read in the +constructor), and then starts: + +* Input indexing worker pool. Input indexing worker pool calls `Index()` on each + chunk source, and then appends the chunk source to the `chunk_sources_` deque + (under mutex). + +* Chunk output worker pool. Fetches the next number from `stream_shuffler_` + under mutex, then reads the chunk from the chunk source using per-source mutex. + +If `stream_shuffler_` runs out of numbers, it's reset to the range +(`last - chunk_window_`, `last`) (and warning message is logged). + +## ShufflingFrameSampler + +The sampler uses reservoir sampling: + +* It has a reservoir of predefined size (1000000 is quite typical) +* Initially it just fills the reservoir with frames from the input queue until + it's full. +* After that, it picks random frames from the reservoir and outputs them, + refilling the used spot from the input queue. +* It closes the output queue when either explicit Close() is called or the input + queue is closed. +* `using FrameType = V6TrainingData;`, use `absl::FixedArray` for + the reservoir. + +## Anchors in ShufflingChunkPool + +The training pipeline aims to start training new epoch when a certain number +of new chunks are available. +To do that, we reset the running counter of chunks in the ShufflingChunkPool +when we start waiting for new chunks, and then wait until the counter reaches +the desired number. +However, when the script restarts, we also want to approximately know how many +new chunks to wait for before starting the training. + +To do that, we use the concept of "anchors". Anchor is just a GetChunkSortKey() +of a given chunk that we remember. + +More specifically, we add the following functions to ShufflingChunkPool: + +* `std::string ResetAnchor()` — resets the anchor to the latest chunk seen so + far, and returns its sort key. Also resets the internal counter of chunks seen + since the anchor. +* `int ChunksSinceAnchor()` — returns the number of chunks seen since the anchor. +* `std::string CurrentAnchor()` — returns the current anchor sort key. +* `void SetAnchor(std::string_view)` — is usually called BEFORE starting processing + chunks. Does not reset the counter, but sets the anchor to the given value. + When the read chunk has the same key as the anchor, the counter is reset to zero. + +Python clients access this functionality by issuing +`StageControlRequest` messages through +`DataLoader.send_control_message()`. The daemon pipeline demonstrates how the +first chunk-pool response is used to update anchor state. + +The anchor functionality works differently during initial load vs. ongoing processing: + +**Initial Load (backward processing):** + +* Chunks are processed in newest-first order during initial scan +* If no anchor is set: count all chunks +* If anchor is set: only count chunks newer than the anchor; reset counter to 0 + when anchor is encountered + +**Ongoing Processing (after initial load):** + +* New chunks are processed as they arrive +* Counter increments by GetChunkCount() for each new chunk source +* Counter resets to 0 when a chunk source matching the anchor is encountered + +Metrics: +In [ShufflingChunkPoolMetricsProto](../proto/training_metrics.proto) we add: + +* `int32 chunks_since_anchor` — number of chunks seen since the anchor. Simple + numerical field. +* `string anchor` — current anchor sort key. + +In [UI](../src/lczero_training/tui/stage_widgets.py) we: + +* Update the last_chunk_key display to have a label "Last:" +* Add a new row "⚓:" for the anchor key. +* Add a new row "Since ⚓:" for chunks_since_anchor. diff --git a/docs/new_stage.md b/docs/new_stage.md new file mode 100644 index 00000000..2ae47e15 --- /dev/null +++ b/docs/new_stage.md @@ -0,0 +1,133 @@ +# Writing a New Data Loader Stage + +This guide walks through the lifecycle of adding another stage to the dynamic +data loader pipeline. It assumes you are working in the C++ orchestrator (under +`csrc/loader/stages/`) and that the Python bindings already consume staged +configurations. + +## 1. Design the Stage Surface + +- **Purpose and data flow**: Decide whether the stage produces new data (no + upstream input) or transforms items from an existing queue. +- **Configuration shape**: Determine which knobs are required during + construction (thread counts, capacities, etc.). These become fields on the + stage-specific protobuf message. +- **Outputs and control hooks**: Clarify what queue type the stage emits and + whether it needs control-plane messages. + +## 2. Extend the Protobufs + +- Add a new `message Config` to `proto/data_loader_config.proto`. +- For single-output stages, add `optional QueueConfig output = N` to configure + the output queue. `QueueConfig` provides `queue_capacity` (default 4), + `overflow_behavior` (BLOCK, DROP_NEW, KEEP_NEWEST), and optional `name`. +- For multi-output stages, use `repeated QueueConfig output` with parallel + configuration arrays (see `ChunkSourceSplitterConfig` for reference). +- Update `StageConfig` with an `optional Config` entry so the stage + can be referenced from the `repeated stage` list. +- If the stage emits custom metrics, extend `StageMetricProto` in + `proto/training_metrics.proto`. Prefer the existing `load_metrics`, + `queue_metrics`, and `count_metrics` collections when possible. +- When the stage needs control requests or responses, extend + `proto/stage_control.proto` so they can be carried through + `StageControlRequest`/`StageControlResponse`. +- Regenerate protobufs (`meson compile -C builddir` or `just build-proto`). + +## 3. Choose a Base Class + +- **Use `SingleInputStage`** when the stage consumes exactly + one upstream queue. The helper provides `input_queue()` to access the typed + `Queue*` and implements `SetInputs()` to wire the input during + initialization. +- **Use `SingleOutputStage`** when the stage produces exactly one + output queue. The helper manages the output queue, implements `GetOutput()` + with name validation, and surfaces the typed `Queue*` via + `output_queue()`. +- **Most stages inherit from both** `SingleInputStage` and `SingleOutputStage` + using virtual inheritance (both base classes virtually inherit from `Stage` + to avoid the diamond problem). Example: + ```cpp + class MyStage : public SingleInputStage, + public SingleOutputStage { + public: + explicit MyStage(const MyStageConfig& config) + : SingleInputStage(config), + SingleOutputStage(config.output()) {} + }; + ``` +- **Inherit `Stage` directly** when the stage has multiple inputs, multiple + outputs, or manages more complex wiring. In that case you must implement + `SetInputs()`, input/output discovery, and `GetOutput()` yourself. +- Place declarations in `csrc/loader/stages/.h` and definitions in + the matching `.cc` file. + +## 4. Implement the Stage API + +- **Constructor**: Initialize base classes with config and `config.output()`. + Store additional config fields and initialize worker pools. + Avoid starting threads here. +- **`SetInputs(absl::Span inputs)`**: Only implement if you + inherit from `Stage` directly. `SingleInputStage` provides this automatically + and validates that exactly one input is provided. For stages with no inputs, + validate that the span is empty. +- **`Start()`**: Launch background work. Acquire `Queue::Producer` instances + from `output_queue()->CreateProducer()` for emitting data and honour + `stop_requested_` flags so shutdown is cooperative. The input queue is + available via `input_queue()` at this point. +- **`Stop()`**: Close queues via `output_queue()->Close()`, signal workers to + exit, and join threads. Remember that downstream stages expect + `Queue::Close()` to signal completion. +- **`GetOutput(std::string_view name)`**: Only implement if the stage has + multiple outputs. `SingleOutputStage` provides this automatically for + single-output stages, including name validation. +- **`Control()`**: Handle relevant `StageControlRequest` sub-messages and return + a populated `StageControlResponse` wrapped in `std::optional`. Return + `std::nullopt` for requests the stage does not recognise. + +## 5. Report Metrics + +- **Accumulate state** while workers run (e.g., load metrics, counters, + queue statistics). +- **`FlushMetrics()`** should snapshot the current values, reset internal + counters as needed, and populate `StageMetricProto`. Use helpers like + `MetricsFromQueue("output", *output_queue())` to expose queue utilisation + under `queue_metrics`, and append load information via `load_metrics`. +- For multiple queues or distinct metric groups, add additional entries with + meaningful names (`"output"`, `"prefetch"`, etc.) so downstream tooling can + pick the right series. +- If you rename or split a metric, document the change and update dashboards. + For example, `ShufflingChunkPool` now emits `chunks_current` (window size) + and `chunks_total` (total indexed chunks) instead of a single `chunks` + series; the Grafana panels consuming the old series were repointed to + `chunks_current` so the graphs remain accurate. + +## 6. Register the Stage + +- Update `CreateStage` in `csrc/loader/stages/stage_factory.cc` to construct the + new class when its config is present. Enforce the "exactly one sub-config" + rule by keeping the existing `CountStageConfigs()` logic in sync. +- Ensure `meson.build` lists the new source files so the static library rebuilds. + +## 7. Wire Up Tests + +- Add focused unit tests under `csrc/loader/stages/` validating constructor + errors, thread lifecycle, metric flushing, and (if applicable) control-plane + behaviour. +- Provide integration coverage where the stage participates in a small pipeline + built from serialized `DataLoaderConfig` messages. +- If Python bindings surface stage-specific behaviour, extend the relevant + `pytest` suites too. + +## 8. Update Documentation and Examples + +- Document new config fields in `docs/` (for example, augment `docs/loader.md` + or create stage-specific notes). +- Add sample snippets or textproto fragments showing how to reference the + stage in a pipeline. +- Mention any new control commands so the daemon/TUI maintainers know how to + surface them. + +Following these steps keeps the stage ecosystem consistent: configurations are +validated at construction time, queues remain type-safe, metrics feed the UI, +and Python clients continue to operate through the generic factory and control +plane. diff --git a/docs/overview.md b/docs/overview.md new file mode 100644 index 00000000..2edac397 --- /dev/null +++ b/docs/overview.md @@ -0,0 +1,13 @@ +# Overview, Glossary, and File Formats + +This document serves as a glossary of terms used in the project and describes +the file formats utilized. + +* Training data **frame** is a data structure that holds information needed for + the NN training about a single chess position. Currently it's a fixed sized + struct, e.g. [V6TrainingData](../libs/lc0/src/trainingdata/trainingdata.h). +* **Chunk** is a sequence of frames from a single game. Currently, they are + stored in a gzipped file where frames are concatenated together. +* **Chunk Source** is a file that contains one or more chunks. In older versions + of the code, it was a single gzipped file. In the new version, it also may be + an uncompressed .tar file. diff --git a/docs/shuffling_pool_hanse_sampling.md b/docs/shuffling_pool_hanse_sampling.md new file mode 100644 index 00000000..ca309f7f --- /dev/null +++ b/docs/shuffling_pool_hanse_sampling.md @@ -0,0 +1,58 @@ +# Implement single position sampling in Shuffling Pool + +This document defines a new way of sampling in +[Shuffling Pool](../csrc/loader/stages/shuffling_chunk_pool.h) (and .cc). + +## reshuffle_count -> use_count + +In `TrainingChunk` in +[training_chunk.h](../csrc/loader/stages/training_chunk.h) the `reshuffle_count` +should be renamed to `use_count`. As with the new sampling method, usage is not +necessarily tied to reshuffling. Also the code that +[uses it](../csrs/loader/stages/chunk_unpacker.cc) will need to be updated. + +In `ChunkSourceItem`, instead of one `reshuffle_count` per entire chunk source, +we'll have a `std::vector use_counts`, initially filled with zeros. +Instead of updating it on reshuffling, we will update it for individual chunks +when they are returned. In `GetNextChunkData()` we return old value (i.e. 0 for +the first time). + +Local struct `ShufflingChunkPool::ChunkData` in +../csrc/loader/stages/shuffling_chunk_pool.cc should also have `use_count` +instead of `reshuffle_count`. + +## Configuration changes + +In the [config](../proto/data_loader_config.proto), in `ShufflingChunkPoolConfig`, +we add the following fields: + +```proto +message ShufflingChunkPoolConfig { + // existing fields... + optional uint64 hanse_sampling_threshold = 6; // by default, do not use new sampling. + optional double hanse_sampling_gamma = 7 [default = 1.0]; +} +``` + +## Algorithm changes + +In addition to `use_counts`, `ChunkSourceItem` will have a new field: +`std::vector num_records;`. It will contain the number of records in +each chunk, and will act as a cache. Initially, it's filled with zeros. + +When `record_bound` is not set, the sampling method is the same as before. + +When `record_bound` is set, we will use the new sampling method. It goes like +this: + +`GetNextChunkData()` doesn't call `LoadChunkData()` right away. +Instead, it first checks if `num_records[chunk_index]` is zero. If so, it +calls `LoadChunkData()` to load the chunk, and counts the number of records in +it (by dividing its size to `sizeof`). + +Then, we decide whether to return this chunk or to sample again. We do this by +drawing a random number `u` uniformly from `[0, 1)`, and comparing it to +`p = min(1.0, num_records/hthreshold) ^ gamma`. If `u < p`, we return this chunk +(we need to call `LoadChunkData()` if we didn't already) and increment +use_count. Otherwise, we sample again (i.e. pick a new `chunk_index` and repeat +the process) — without incrementing `use_count`. \ No newline at end of file diff --git a/docs/training_tuple.md b/docs/training_tuple.md new file mode 100644 index 00000000..aa0bc57c --- /dev/null +++ b/docs/training_tuple.md @@ -0,0 +1,31 @@ +# Training Tuple Format + +The `convert_v6_to_tuple` function in `tf/chunkparser.py` processes training +data and produces a 5-element tuple: +`(planes, probs, winner, best_q, plies_left)`. + +When these raw byte strings are interpreted as NumPy arrays, they have the +following shapes for each training example: + +1. **`planes`**: `(112, 64)` as a `float32` array. + * This represents the board state as 112 feature planes, each of size 8x8 + (64). The original 104 planes from the input are augmented with 8 + additional planes for information like castling rights, side to move, + the rule 50 count, and board edge detection. + +2. **`probs`**: `(1858,)` as a `float32` array. + * This corresponds to the `float probabilities[1858]` member in the + `V6TrainingData` C++ struct, representing the policy probabilities for + all possible moves. + +3. **`winner`**: `(3,)` as a `float32` array. + * This holds the game's outcome from the current player's perspective, + representing the probabilities for a win, draw, and loss, respectively. + +4. **`best_q`**: `(3,)` as a `float32` array. + * Similar to `winner`, this stores the value of the position after search + (the Q-value), also represented as win, draw, and loss probabilities. + +5. **`plies_left`**: A scalar `float32`. + * This value represents the estimated number of plies remaining until the + end of the game. diff --git a/docs/tui.md b/docs/tui.md new file mode 100644 index 00000000..2963cebb --- /dev/null +++ b/docs/tui.md @@ -0,0 +1,124 @@ +# UI Design + +The application will present a single-screen dashboard with a classic blue background, organized into several key, always-visible panes. + +## 1. Overall Layout + +The screen is divided into four main sections: + +1. **Header Bar (Top):** A slim, single-line bar at the very top. +2. **Data Pipeline Pane (Top-Left/Main):** The largest and most detailed pane. +3. **JAX Training Status Pane (Right):** A pane dedicated to live training metrics. +4. **Log Pane (Bottom):** A pane for displaying raw log output. + +## 2. Header Bar + +This bar provides high-level, global status at a glance. + +- **Uptime:** Total wall-clock time since the script was launched. +- **Overall Stage:** The current high-level state of the application. Will display one of: `WAITING FOR DATA`, `TRAINING`, `EXPORTING`, or `ERROR`. + +### 3. Data Pipeline Pane + +This pane visualizes the flow of data through the C++ pipeline. The new +design is a vertically scrollable list where every stage and queue consumes +one row (two only when extra detail is needed). Rows are rendered in the same +order as traffic flows through the loader, preserving the mental model of the +pipeline without relying on borders or grid positioning. + +- **Pipeline Stages (Rows):** Each stage of the C++ pipeline is rendered as a + single line. + - **Stage Heading:** The label uses the canonical stage name reported in the + metrics, so newly added stages automatically appear without additional UI + wiring. + - **Load Metrics:** Stages with thread pools render load in `load + active/total` format. + - **Specific Stats:** Each metric is rendered as an individual "chip" widget, + e.g. skipped file counters or pool sizes, so additional detail can be added + without breaking alignment. + +- **Queues (Rows):** Each stage row is followed by a queue row that surfaces the + metrics of the outgoing queue. + - **Queue Heading:** The row label is the canonical stage name reported by the + daemon. The first chip shows the queue name when it differs from the + default. + - **Throughput:** A chip displays the 1-second `items/s` rate and turns red + when the rate drops to zero. + - **Totals:** A chip shows the lifetime count of elements that passed through + the queue, formatted with apostrophes. + - **Fill State:** A horizontal progress bar visualises the average queue fill + against capacity, followed by a chip with the numeric `avg/capacity` + display. Unknown values fall back to `--`. + +- **Train/Validation/Test Splitter:** + - The pipeline view will show a "Stream Splitter" stage. + - Hotkeys (e.g., F1, F2, F3) will allow the user to instantly switch the view + - After this stage, the UI will display the pipeline stats for **one** stream + at a time (defaulting to 'Training'). to show the stats for the 'Training', + 'Validation', or 'Test' streams. + +### 4. Training schedule/pipeline pane + +The area below the Data Pipeline Pane is split horizontally into two +sections: Training Schedule (left) and JAX Training Status (right). + +The Training Schedule pane shows: + +- **Combined uptime and stage line**: "Uptime: 2d 14:30:45 Stage: TRAINING" + - Uptime includes days when >24 hours (format: "2d 14:30:45") + - Stage shows current training state (WAITING_FOR_DATA, TRAINING, EXPORTING, ERROR) +- **Completed epochs**: Simple counter of epochs completed since daemon start +- **New chunks progress bar**: Shows chunks collected since training start vs. target + - Indeterminate state when target is unknown (0) +- **Training time progress bar**: Current training time vs. previous training duration + - Indeterminate state when no previous duration exists +- **Cycle time progress bar**: Current cycle time vs. previous cycle duration + - Indeterminate state when no previous duration exists + +**Implementation details:** + +- **Header bar**: Completely empty (no content) +- **Data structure**: `TrainingScheduleData` dataclass with all timing fields: + - `current_stage: TrainingStage` (enum) + - `completed_epochs_since_start: int` + - `new_chunks_since_training_start: int` + - `chunks_to_wait: int` + - `total_uptime_seconds: float` + - `current_training_time_seconds: float` + - `previous_training_time_seconds: float` + - `current_cycle_time_seconds: float` + - `previous_cycle_time_seconds: float` +- **Timing computation**: All timing values computed in daemon, not TUI +- **Progress bars**: Show indeterminate state when maximum values ≤ 0 +- **Layout**: Compact single-line widgets with no extra padding/margins +- **Files**: + - `training_widgets.py`: Widget implementations + - `pipeline.py`: Training state tracking and data collection + - `daemon.py`: Metrics collection and transmission + - `messages.py`: Protocol definitions with enum serialization support + +### 5. JAX Training Status Pane + +This pane is dedicated to the live status of an active JAX training run. It +remains blank or shows summary info when the system is not actively training. + +- **Epoch Progress:** A progress bar showing completion of the current epoch + (`Step 12345 / 50000`). +- **Performance Metrics:** + - Steps per second: `345.6 steps/s`. + - Estimated Time Remaining (ETR) for the current run. + - Total wall time spent on the current training run. +- **Loss Values (Numerical Only):** + - A prominent display of the **Total Loss**. + - A compact 2-column grid displaying the individual values for the **7 Head + Losses**. + +### 6. Log Pane + +A pane across the bottom of the screen. + +- **Content:** A direct feed of all output sent to `stderr` from any part of the + application (Python or C++). +- **Functionality:** The pane will be scrollable and will hold a fixed number of + lines (e.g., 1000) to prevent unbounded memory usage, discarding the oldest + lines as new ones arrive. diff --git a/docs/weights_tool.md b/docs/weights_tool.md new file mode 100644 index 00000000..e3fc0518 --- /dev/null +++ b/docs/weights_tool.md @@ -0,0 +1,396 @@ +# lc0-weights - Weight Manipulation Tool + +## Overview + +`lc0-weights` is a command-line tool and Python library for manipulating Leela Chess Zero neural network weight files. It enables arithmetic operations on networks, component grafting, and format conversion without requiring JAX or other heavy dependencies. + +**Key capabilities:** + +* **Arithmetic operations**: Add, subtract, multiply networks (e.g., model interpolation/averaging) +* **Grafting**: Replace specific network components (policy heads, value heads, encoder layers) +* **Format conversion**: Convert between LINEAR16, FLOAT16, and BFLOAT16 encodings +* **Pure Python/NumPy**: No JAX or TensorFlow dependencies required + +## Quick Start + +Interpolate two networks with equal weights: + +```bash +# Using inline expression +uv run lc0-weights \ + --expr "output = weights('network_a.pb.gz') * 0.5 + weights('network_b.pb.gz') * 0.5" \ + --output interpolated.pb.gz + +# Using a script file +echo "output = weights('network_a.pb.gz') * 0.5 + weights('network_b.pb.gz') * 0.5" > interpolate.py +uv run lc0-weights interpolate.py --output interpolated.pb.gz + +# Using stdin +echo "output = weights('network_a.pb.gz') * 0.5 + weights('network_b.pb.gz') * 0.5" | \ + uv run lc0-weights --output interpolated.pb.gz +``` + +## Command-Line Interface + +### Arguments + +| Argument | Required | Description | +| ------------ | -------- | -------------------------------------------------------------------- | +| `--expr` | No | Python expression to execute | +| `script` | No | Path to Python script file (positional, used if --expr not given) | +| `--input` | No | Pre-load input as `NAME=PATH` (can be used multiple times) | +| `--output` | No | Output path (if `output` variable is set in expression) | +| `--encoding` | No | Output encoding format: LINEAR16, FLOAT16 (default), or BFLOAT16 | + +**Note:** You must provide either `--expr`, a script file path, or pipe input via stdin. + +### Available Functions in --expr + +When executing expressions with `--expr`, the following are available: + +* `weights(path)`: Load a weight file +* `save(net, path, encoding='FLOAT16')`: Save a network +* `np`: NumPy module +* `lc0`: Protobuf module (for accessing constants) + +### CLI Examples + +#### Simple Interpolation + +**Using inline expression:** +```bash +uv run lc0-weights \ + --expr "output = weights('A.pb.gz') * 0.5 + weights('B.pb.gz') * 0.5" \ + --output result.pb.gz +``` + +**Using a script file:** +```bash +# Create script file +cat > interpolate.py << 'EOF' +A = weights('A.pb.gz') +B = weights('B.pb.gz') +output = A * 0.5 + B * 0.5 +EOF + +# Run script +uv run lc0-weights interpolate.py --output result.pb.gz +``` + +**Using stdin:** +```bash +echo "output = weights('A.pb.gz') * 0.5 + weights('B.pb.gz') * 0.5" | \ + uv run lc0-weights --output result.pb.gz +``` + +#### Using Input Aliases + +Pre-load networks to simplify expressions: + +```bash +uv run lc0-weights \ + --input A=network_a.pb.gz \ + --input B=network_b.pb.gz \ + --expr "output = A * 0.9 + B * 0.1" \ + --output result.pb.gz +``` + +#### Grafting a Policy Head + +Replace the policy head of one network with another: + +**Using inline expression:** +```bash +uv run lc0-weights --expr " +base = weights('base_network.pb.gz') +donor = weights('network_with_better_policy.pb.gz') +base.weights.policy = donor.weights.policy +base.save('grafted.pb.gz') +" +``` + +**Using a script file (recommended for multi-line operations):** +```bash +# Create script +cat > graft_policy.py << 'EOF' +base = weights('base_network.pb.gz') +donor = weights('network_with_better_policy.pb.gz') +base.weights.policy = donor.weights.policy +base.save('grafted.pb.gz') +EOF + +# Run script +uv run lc0-weights graft_policy.py +``` + +#### Format Conversion + +Convert a network to BFLOAT16 encoding: + +```bash +uv run lc0-weights \ + --input net=network.pb.gz \ + --expr "output = net" \ + --output converted.pb.gz \ + --encoding BFLOAT16 +``` + +#### Complex Expression + +Weighted average with custom formula: + +```bash +uv run lc0-weights \ + --input A=net1.pb.gz \ + --input B=net2.pb.gz \ + --input C=net3.pb.gz \ + --expr "output = A * 0.5 + B * 0.3 + C * 0.2" \ + --output averaged.pb.gz +``` + +## Python Library Usage + +### Importing + +```python +from lczero_training.tools import load_weights, save_weights +``` + +### Loading and Saving Weights + +```python +# Load a network +net = load_weights("network.pb.gz") + +# Save with different encoding +save_weights(net, "output.pb.gz", encoding="FLOAT16") +``` + +### Arithmetic Operations + +```python +# Load networks +net_a = load_weights("network_a.pb.gz") +net_b = load_weights("network_b.pb.gz") + +# Interpolation (model averaging) +interpolated = net_a * 0.7 + net_b * 0.3 + +# Addition +combined = net_a + net_b + +# Subtraction +difference = net_a - net_b + +# Scalar multiplication +scaled = net_a * 0.5 + +# Save result +save_weights(interpolated, "result.pb.gz") +``` + +### Accessing and Modifying Components + +```python +net = load_weights("network.pb.gz") + +# Access nested weight arrays +q_weights = net.weights.encoder[0].mha.q_w.value # Returns NumPy array +print(q_weights.shape) + +# Modify weights +net.weights.encoder[0].mha.q_w.value = q_weights * 1.1 + +# Save modified network +save_weights(net, "modified.pb.gz") +``` + +### Grafting Components + +```python +base = load_weights("base.pb.gz") +donor = load_weights("donor.pb.gz") + +# Replace policy head +base.weights.policy = donor.weights.policy + +# Replace value head +base.weights.value_heads = donor.weights.value_heads + +# Replace specific encoder layer +base.weights.encoder[0] = donor.weights.encoder[0] + +save_weights(base, "grafted.pb.gz") +``` + +## Weight Encoding Formats + +The tool supports three encoding formats for weight storage: + +| Format | Description | Precision | File Size | +| -------- | -------------------------------------------------------------------------------- | --------- | --------- | +| LINEAR16 | Quantized 16-bit integer with min/max range. Default for Lc0. | ~4 digits | Smallest | +| FLOAT16 | Native IEEE 754 half-precision floating point. | ~3 digits | Medium | +| BFLOAT16 | Brain float 16 (truncated float32). Better range than FLOAT16, less precision. | ~2 digits | Medium | + +**Notes:** + +* LINEAR16 provides good compression with acceptable precision for neural networks +* FLOAT16 is the default for this tool (good balance of precision and size) +* BFLOAT16 is useful when training range matters more than mantissa precision +* All formats are converted to float32 when loaded for arithmetic operations + +## Common Use Cases + +### Network Interpolation (Model Averaging) + +Combine two networks to create a smoother model or blend different training runs: + +```python +from lczero_training.tools import load_weights, save_weights + +net1 = load_weights("run1_final.pb.gz") +net2 = load_weights("run2_final.pb.gz") + +# Average the networks +averaged = net1 * 0.5 + net2 * 0.5 + +save_weights(averaged, "averaged_network.pb.gz") +``` + +### Exponential Moving Average (EMA) + +Update a running average network with a new checkpoint: + +```python +ema_net = load_weights("ema.pb.gz") +new_net = load_weights("latest_checkpoint.pb.gz") + +# EMA with decay 0.999 +ema_updated = ema_net * 0.999 + new_net * 0.001 + +save_weights(ema_updated, "ema.pb.gz") +``` + +### Policy Head Replacement + +Replace a network's policy head (useful for policy distillation): + +```python +student = load_weights("student_network.pb.gz") +teacher = load_weights("teacher_network.pb.gz") + +# Replace student's policy head with teacher's +student.weights.policy_heads = teacher.weights.policy_heads + +save_weights(student, "student_with_teacher_policy.pb.gz") +``` + +### Extracting Network Statistics + +```python +net = load_weights("network.pb.gz") + +# Get statistics from first encoder layer +layer = net.weights.encoder[0].mha.q_w.value +print(f"Shape: {layer.shape}") +print(f"Mean: {layer.mean():.6f}") +print(f"Std: {layer.std():.6f}") +print(f"Min: {layer.min():.6f}") +print(f"Max: {layer.max():.6f}") +``` + +### Format Conversion for Size Optimization + +```python +from lczero_training.tools import load_weights, save_weights + +# Load network (any format) +net = load_weights("large_network.pb.gz") + +# Save with more aggressive compression +save_weights(net, "compressed_network.pb.gz", encoding="LINEAR16") +``` + +## Advanced Usage + +### Accessing Nested Structures + +The weight wrapper provides Pythonic access to the nested protobuf structure: + +```python +net = load_weights("network.pb.gz") + +# Access input embedding weights +input_weights = net.weights.input.weights.value + +# Access specific encoder layer +encoder_layer_5 = net.weights.encoder[5] + +# Access multi-head attention components +q_weights = net.weights.encoder[0].mha.q_w.value +k_weights = net.weights.encoder[0].mha.k_w.value +v_weights = net.weights.encoder[0].mha.v_w.value + +# Access policy head +policy_weights = net.weights.policy_heads.vanilla.ip_pol_w.value +``` + +### Complex Weighted Combinations + +```python +nets = [load_weights(f"checkpoint_{i}.pb.gz") for i in range(5)] +weights = [0.1, 0.15, 0.2, 0.25, 0.3] # More weight to recent checkpoints + +result = sum(net * w for net, w in zip(nets, weights)) +save_weights(result, "weighted_ensemble.pb.gz") +``` + +### Selective Component Grafting + +```python +base = load_weights("base.pb.gz") +donor = load_weights("donor.pb.gz") + +# Replace only the first 10 encoder layers +for i in range(10): + base.weights.encoder[i] = donor.weights.encoder[i] + +save_weights(base, "partial_graft.pb.gz") +``` + +### Working with NumPy Arrays Directly + +All weight access returns NumPy arrays, allowing arbitrary transformations: + +```python +net = load_weights("network.pb.gz") + +# Get layer weights +layer_weights = net.weights.encoder[0].mha.q_w.value + +# Apply custom transformation +import numpy as np +layer_weights_normalized = layer_weights / np.linalg.norm(layer_weights, axis=-1, keepdims=True) + +# Write back +net.weights.encoder[0].mha.q_w.value = layer_weights_normalized + +save_weights(net, "normalized.pb.gz") +``` + +## Implementation Details + +### Lazy Loading + +Weights are decoded from their compressed format only when accessed, and cached for subsequent use. This makes the tool memory-efficient when working with large networks. + +### File Format Support + +Both `.pb` (uncompressed protobuf) and `.pb.gz` (gzip-compressed) formats are supported. The tool automatically detects the format based on the file extension. + +### Arithmetic Semantics + +* Operations are element-wise across all matching layers +* Networks must have compatible structures (same number of encoder layers, etc.) +* Results are computed in float32 precision regardless of input encoding diff --git a/init.sh b/init.sh index 79970b22..bd56f2c2 100755 --- a/init.sh +++ b/init.sh @@ -1,5 +1,4 @@ #!/usr/bin/env bash -protoc --proto_path=libs/lczero-common --python_out=tf libs/lczero-common/proto/net.proto -protoc --proto_path=libs/lczero-common --python_out=tf libs/lczero-common/proto/chunk.proto +protoc --proto_path=libs/lc0 --python_out=tf proto/net.proto touch tf/proto/__init__.py diff --git a/justfile b/justfile new file mode 100644 index 00000000..9b4d7c21 --- /dev/null +++ b/justfile @@ -0,0 +1,84 @@ +# List available commands +default: + @just --list + +# Check if all C++ files in csrc/ are formatted according to clang-format +check-cpp: + find csrc/ -name "*.cpp" -o -name "*.cc" -o -name "*.cxx" -o -name "*.h" -o -name "*.hpp" | xargs clang-format --dry-run --Werror + +# Format all C++ files in csrc/ using clang-format +format-cpp: + find csrc/ -name "*.cpp" -o -name "*.cc" -o -name "*.cxx" -o -name "*.h" -o -name "*.hpp" | xargs clang-format -i + +# Check if all protobuf files are formatted according to clang-format +check-proto: + find proto/ -name "*.proto" | xargs clang-format --dry-run --Werror + +# Format all protobuf files using clang-format +format-proto: + find proto/ -name "*.proto" | xargs clang-format -i + +# Build Python protobuf files +build-proto: + mkdir -p src/proto + touch src/proto/__init__.py + uv run python -m grpc_tools.protoc \ + --proto_path=. \ + --proto_path=libs/lc0 \ + --python_out=src/ \ + --pyi_out=src/ \ + proto/*.proto + uv run python -m grpc_tools.protoc \ + --proto_path=. \ + --proto_path=libs/lc0 \ + --python_out=src/ \ + --pyi_out=src/ \ + proto/net.proto \ + proto/onnx.proto \ + proto/hlo.proto + +# Check if all Python files in src/ are formatted according to ruff +check-python: + uv run ruff check src/ + uv run ruff check --select I src/ + uv run ruff format --check src/ + uv run mypy -p lczero_training --disallow-untyped-defs --disallow-incomplete-defs + +# Format all Python files in src/ using ruff +format-python: + uv run ruff check --fix --select I src/ + uv run ruff format src/ + uv run ruff check --fix src/ + +format: format-cpp format-proto format-python + +# Setup meson build directory with clang +setup-build: + CXX=clang++ CC=clang uv run meson setup build/release \ + --buildtype=release --native-file=native.ini + +# Build the project +build: + uv run meson compile -C build/release/ + +# Create symlink for the built extension module +mksymlink: + cd src/lczero_training && ln -sfT ../../build/release/_lczero_training.cpython-*-x86_64-linux-gnu.so _lczero_training.so + +# Delete build/release, re-setup, rebuild, and re-link +rebuild: && setup-build build mksymlink + rm -rf build/release + +# Run tests +test-cpp: + uv run meson test -C build/release/ + +test-python: + uv run pytest + +test: test-cpp test-python + +check: check-cpp check-proto check-python + +# Run all checks (formatting, build, and tests) +pre-commit: build-proto check build test diff --git a/libs/lc0 b/libs/lc0 new file mode 160000 index 00000000..39557383 --- /dev/null +++ b/libs/lc0 @@ -0,0 +1 @@ +Subproject commit 395573837488074361fad72096bc6feca345d646 diff --git a/libs/lczero-common b/libs/lczero-common deleted file mode 160000 index 00fd892e..00000000 --- a/libs/lczero-common +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 00fd892e648160c294346c87449126d9bad80a16 diff --git a/meson.build b/meson.build new file mode 100644 index 00000000..b970a1b6 --- /dev/null +++ b/meson.build @@ -0,0 +1,401 @@ +project( + 'lczero-training', + 'cpp', + version : '0.1', + meson_version : '>= 1.3.0', + default_options : [ + 'warning_level=3', + 'cpp_std=c++20', + + ], +) + +# Allow Clang nullability extensions when using Clang compiler +cpp_compiler = meson.get_compiler('cpp') +if cpp_compiler.get_id() == 'clang' + add_project_arguments('-Wno-nullability-extension', language : 'cpp') +endif + + + +# External dependencies +zlib_dep = dependency('zlib') + +# Python and PyBind11 dependencies for Python extension +python3 = import('python').find_installation() +pybind11_dep = dependency('pybind11') + +# Abseil dependencies always resolved from the wrap subproject. +absl_proj = subproject('abseil-cpp') +absl_dep_specs = [ + ['log', 'absl_log_dep'], + ['log_initialize', 'absl_log_dep'], + ['check', 'absl_log_dep'], + ['hash', 'absl_hash_dep'], + ['raw_hash_set', 'absl_container_dep'], + ['synchronization', 'absl_synchronization_dep'], + ['random_random', 'absl_random_dep'], + ['flags', 'absl_flags_dep'], + ['flags_parse', 'absl_flags_dep'], + ['throw_delegate', 'absl_base_dep'], +] +absl_deps = {} +foreach spec : absl_dep_specs + absl_deps += { + spec[0].underscorify() : absl_proj.get_variable(spec[1]).as_system() + } +endforeach + +# Test dependencies +gtest_dep = dependency('gtest').as_system() +gtest_main_dep = dependency('gtest_main').as_system() + +# Gaviota tablebase subproject +gaviota_dep = subproject('gaviotatb').get_variable('gaviotatb_dep') + +# Common dependency sets +external_deps = [zlib_dep] +core_absl_deps = [ + absl_deps['log'], + absl_deps['check'], + absl_deps['hash'], + absl_deps['raw_hash_set'], + absl_deps['synchronization'], + absl_deps['random_random'], +] +loader_deps = external_deps + core_absl_deps + [gaviota_dep] +main_deps = loader_deps + [absl_deps['log_initialize']] +# test_deps will be defined after proto_dep +cli_deps = [ + absl_deps['log'], + absl_deps['log_initialize'], + absl_deps['flags'], + absl_deps['flags_parse'], +] + +# Protobuf compilation setup +compile_proto = find_program('libs/lc0/scripts/compile_proto.py') +proto_gen = generator(compile_proto, output: ['@BASENAME@.pb.h'], + arguments : [ + '--proto_path=@CURRENT_SOURCE_DIR@', + '--cpp_out=@BUILD_DIR@', + '@INPUT@']) + +lc0_proto_gen = generator(compile_proto, output: ['@BASENAME@.pb.h'], + arguments : [ + '--proto_path=' + join_paths(meson.current_source_dir(), 'libs/lc0'), + '--cpp_out=@BUILD_DIR@', + '@INPUT@']) + +includes = include_directories('csrc', 'libs/lc0/src') + +add_project_arguments('-DNO_PEXT', language : 'cpp') + +rescorer_files = [ + 'libs/lc0/src/chess/board.cc', + 'libs/lc0/src/chess/gamestate.cc', + 'libs/lc0/src/chess/position.cc', + 'libs/lc0/src/neural/decoder.cc', + 'libs/lc0/src/neural/encoder.cc', + 'libs/lc0/src/trainingdata/reader.cc', + 'libs/lc0/src/trainingdata/trainingdata.cc', + 'libs/lc0/src/trainingdata/writer.cc', + 'libs/lc0/src/utils/commandline.cc', + 'libs/lc0/src/utils/configfile.cc', + 'libs/lc0/src/utils/esc_codes.cc', + 'libs/lc0/src/utils/optionsdict.cc', + 'libs/lc0/src/utils/optionsparser.cc', + 'libs/lc0/src/utils/random.cc', + 'libs/lc0/src/utils/string.cc', +] + +if host_machine.system() == 'windows' + rescorer_files += 'libs/lc0/src/utils/filesystem.win32.cc' +else + rescorer_files += 'libs/lc0/src/utils/filesystem.posix.cc' +endif + +files = [ + 'csrc/loader/chunk_source/debug_chunk_source.cc', + 'csrc/loader/chunk_source/rawfile_chunk_source.cc', + 'csrc/loader/chunk_source/tar_chunk_source.cc', + 'csrc/loader/data_loader_metrics.cc', + 'csrc/loader/data_loader.cc', + 'csrc/loader/stages/chunk_rescorer.cc', + 'csrc/loader/stages/chunk_source_loader.cc', + 'csrc/loader/stages/chunk_source_splitter.cc', + 'csrc/loader/stages/chunk_unpacker.cc', + 'csrc/loader/stages/file_path_provider.cc', + 'csrc/loader/stages/join_stage.cc', + 'csrc/loader/stages/position_sampling.cc', + 'csrc/loader/stages/position_sampling.cc', + 'csrc/loader/stages/shuffling_chunk_pool.cc', + 'csrc/loader/stages/shuffling_frame_sampler.cc', + 'csrc/loader/stages/simple_chunk_extractor.cc', + 'csrc/loader/stages/stage_factory.cc', + 'csrc/loader/stages/stage.cc', + 'csrc/loader/stages/tensor_generator.cc', + 'csrc/utils/gz.cc', + 'csrc/utils/stream_shuffler.cc', + 'csrc/utils/training_data_printer.cc', + 'libs/lc0/src/syzygy/syzygy.cc', + 'libs/lc0/src/trainingdata/rescorer.cc', + 'libs/lc0/src/utils/files.cc', + 'libs/lc0/src/utils/logging.cc', + 'libs/lc0/src/utils/protomessage.cc', +] + rescorer_files + +# Process protobuf files for C++ +proto_files = [ + proto_gen.process('proto/data_loader_config.proto', + preserve_path_from : meson.current_source_dir()), + proto_gen.process('proto/training_metrics.proto', + preserve_path_from : meson.current_source_dir()), + proto_gen.process('proto/stage_control.proto', + preserve_path_from : meson.current_source_dir()), + lc0_proto_gen.process('libs/lc0/proto/net.proto', + preserve_path_from : join_paths(meson.current_source_dir(), 'libs/lc0')), + lc0_proto_gen.process('libs/lc0/proto/onnx.proto', + preserve_path_from : join_paths(meson.current_source_dir(), 'libs/lc0')), + lc0_proto_gen.process('libs/lc0/proto/hlo.proto', + preserve_path_from : join_paths(meson.current_source_dir(), 'libs/lc0')) +] + +# Create a dependency for protobuf files that tests can use +proto_dep = declare_dependency(sources: proto_files) +test_deps = [gtest_dep, gtest_main_dep, proto_dep] + + +files += proto_files + +loader_lib = static_library( + 'loader', + files, + include_directories : includes, + dependencies : loader_deps, +) + +exe = executable( + 'loader', + 'csrc/loader/loader_main.cpp', + include_directories : includes, + dependencies : cli_deps + [proto_dep], + link_with : loader_lib, +) + +stream_shuffler_test = executable( + 'stream_shuffler_test', + 'csrc/utils/stream_shuffler_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['random_random']], + link_with : loader_lib, +) + +queue_test = executable( + 'queue_test', + 'csrc/utils/queue_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], +) + +file_path_provider_test = executable( + 'file_path_provider_test', + 'csrc/loader/stages/file_path_provider_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], + link_with : loader_lib, +) + +chunk_source_loader_test = executable( + 'chunk_source_loader_test', + 'csrc/loader/stages/chunk_source_loader_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization']], + link_with : loader_lib, +) + +shuffling_chunk_pool_test = executable( + 'shuffling_chunk_pool_test', + 'csrc/loader/stages/shuffling_chunk_pool_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], + link_with : loader_lib, +) + +# simple_chunk_extractor_test = executable( +# 'simple_chunk_extractor_test', +# 'csrc/loader/stages/simple_chunk_extractor_test.cc', +# include_directories : includes, +# dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], +# link_with : loader_lib, +# ) + +chunk_rescorer_test = executable( + 'chunk_rescorer_test', + 'csrc/loader/stages/chunk_rescorer_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], + link_with : loader_lib, +) + +chunk_unpacker_test = executable( + 'chunk_unpacker_test', + 'csrc/loader/stages/chunk_unpacker_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], + link_with : loader_lib, +) + +shuffling_frame_sampler_test = executable( + 'shuffling_frame_sampler_test', + 'csrc/loader/stages/shuffling_frame_sampler_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['random_random']], + link_with : loader_lib, +) + +tensor_test = executable( + 'tensor_test', + 'csrc/utils/tensor_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['throw_delegate']], +) + +tensor_generator_test = executable( + 'tensor_generator_test', + 'csrc/loader/stages/tensor_generator_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization']], + link_with : loader_lib, +) + +stats_test = executable( + 'stats_test', + 'csrc/utils/metrics/stats_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization']], + link_with : loader_lib, +) + +load_metric_test = executable( + 'load_metric_test', + 'csrc/utils/metrics/load_metric_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization']], + link_with : loader_lib, +) + +stage_factory_test = executable( + 'stage_factory_test', + 'csrc/loader/stages/stage_factory_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], + link_with : loader_lib, +) + +data_loader_test = executable( + 'data_loader_test', + 'csrc/loader/data_loader_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], + link_with : loader_lib, +) + +join_stage_test = executable( + 'join_stage_test', + 'csrc/loader/stages/join_stage_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], + link_with : loader_lib, +) +test('stream_shuffler_test', stream_shuffler_test) +test('queue_test', queue_test) +test('file_path_provider_test', file_path_provider_test) +test('chunk_source_loader_test', chunk_source_loader_test) +chunk_source_splitter_test = executable( + 'chunk_source_splitter_test', + 'csrc/loader/stages/chunk_source_splitter_test.cc', + include_directories : includes, + dependencies : test_deps + [absl_deps['synchronization'], absl_deps['log']], + link_with : loader_lib, +) +test('chunk_source_splitter_test', chunk_source_splitter_test) +test('shuffling_chunk_pool_test', shuffling_chunk_pool_test) +# test('simple_chunk_extractor_test', simple_chunk_extractor_test) +test('chunk_rescorer_test', chunk_rescorer_test) +test('chunk_unpacker_test', chunk_unpacker_test) +test('shuffling_frame_sampler_test', shuffling_frame_sampler_test) +test('tensor_test', tensor_test) +test('tensor_generator_test', tensor_generator_test) +test('stats_test', stats_test) +test('load_metric_test', load_metric_test) +test('stage_factory_test', stage_factory_test) +test('data_loader_test', data_loader_test) +test('join_stage_test', join_stage_test) + +file_path_provider_main = executable( + 'file_path_provider_main', + 'csrc/loader/stages/file_path_provider_main.cc', + include_directories : includes, + dependencies : cli_deps + [proto_dep], + link_with : loader_lib, +) + +dump_chunk = executable( + 'dump_chunk', + 'csrc/tools/dump_chunk_main.cc', + include_directories : includes, + dependencies : cli_deps + [zlib_dep], + link_with : loader_lib, +) + +filter_chunks = executable( + 'filter_chunks', + 'csrc/tools/filter_chunks_main.cc', + include_directories : includes, + dependencies : cli_deps + [proto_dep, zlib_dep], + link_with : loader_lib, +) + +result_distribution = executable( + 'result_distribution', + 'csrc/tools/result_distribution_main.cc', + include_directories : includes, + dependencies : cli_deps + [proto_dep, absl_deps['synchronization']], + link_with : loader_lib, +) + +startpos_policy_distribution = executable( + 'startpos_policy_distribution', + 'csrc/tools/startpos_policy_distribution_main.cc', + include_directories : includes, + dependencies : cli_deps + [proto_dep], + link_with : loader_lib, +) + +rescore_chunk = executable( + 'rescore_chunk', + 'csrc/tools/rescore_chunk_main.cc', + include_directories : includes, + dependencies : cli_deps + [proto_dep], + link_with : loader_lib, +) + +position_weight_stats = executable( + 'position_weight_stats', + 'csrc/tools/position_weight_stats_main.cc', + include_directories : includes, + dependencies : cli_deps + [proto_dep], + link_with : loader_lib, +) + +# Python extension module +python3.extension_module( + '_lczero_training', + 'csrc/loader/pybind_module.cc', + include_directories : includes, + dependencies : [pybind11_dep, proto_dep, absl_deps['log_initialize']] + loader_deps, + link_with : loader_lib, + install : true, + subdir : 'lczero_training', +) diff --git a/native.ini b/native.ini new file mode 100644 index 00000000..cd8a628e --- /dev/null +++ b/native.ini @@ -0,0 +1,2 @@ +[binaries] +python = '@GLOBAL_SOURCE_ROOT@/.venv/bin/python' \ No newline at end of file diff --git a/proto/checkpoint_migration_config.proto b/proto/checkpoint_migration_config.proto new file mode 100644 index 00000000..e9046869 --- /dev/null +++ b/proto/checkpoint_migration_config.proto @@ -0,0 +1,15 @@ +syntax = "proto3"; + +package lczero_training.proto; + +message CheckpointMigrationRule { + // Path in the old state pytree. + string from_path = 1; + + // Path in the new state pytree. + string to_path = 2; +} + +message CheckpointMigrationConfig { + repeated CheckpointMigrationRule rule = 1; +} diff --git a/proto/data_loader_config.proto b/proto/data_loader_config.proto new file mode 100644 index 00000000..dd0e4018 --- /dev/null +++ b/proto/data_loader_config.proto @@ -0,0 +1,206 @@ +syntax = "proto2"; + +package lczero.training; + +// Configuration for output queue used by stages. +message QueueConfig { + // Queue overflow behavior. + enum OverflowBehavior { + BLOCK = 0; + DROP_NEW = 1; + KEEP_NEWEST = 2; + } + + // Optional name for the output (used for multi-output stages). + optional string name = 1; + // Size of the output queue. + optional uint64 queue_capacity = 2 [default = 4]; + // Overflow behavior of the output queue. + optional OverflowBehavior overflow_behavior = 3; +} + +// Configuration for file path provider that watches directories for new +// training files. Maps to FilePathProviderOptions in +// csrc/loader/chunk_feed/file_path_provider.h +message FilePathProviderConfig { + // Path to directory containing training data files. + optional string directory = 1; + // Output queue configuration. + optional QueueConfig output = 2; +} + +// Configuration for chunk source loader that converts file paths to chunk +// sources. Maps to ChunkSourceLoaderOptions in +// csrc/loader/chunk_feed/chunk_source_loader.h +message ChunkSourceLoaderConfig { + enum FrameFormat { + V6TrainingData = 0; + V7TrainingData = 1; + } + // Number of worker threads for loading. + optional uint64 threads = 1 [default = 1]; + // Output queue configuration. + optional QueueConfig output = 2; + // Training data frame format. + optional FrameFormat frame_format = 3; +} + +message PositionSamplingConfig { + optional float diff_focus_q_weight = 1 [default = 0.0]; + optional float diff_focus_pol_scale = 2 [default = 1.0]; + optional float diff_focus_alpha = 3 [default = 1.0]; + optional float diff_focus_beta = 4 [default = 0.0]; + optional float diff_focus_gamma = 5 [default = 1.0]; + optional float diff_focus_tau = 6 [default = 1.0]; + optional float default_weight = 7 [default = 1.0]; +} + +// Configuration for shuffling chunk pool that manages chunk shuffling and +// loading. Maps to ShufflingChunkPoolOptions in +// csrc/loader/chunk_feed/shuffling_chunk_pool.h +message ShufflingChunkPoolConfig { + // Size of the chunk shuffle buffer. + optional uint64 chunk_pool_size = 1 [default = 100000]; + // Threads for ingesting new chunk sources (lightweight). + optional uint64 source_ingestion_threads = 3 [default = 1]; + // Threads for loading chunk data. + optional uint64 chunk_loading_threads = 4 [default = 4]; + // Output queue configuration. + optional QueueConfig output = 5; + // When set to a positive value, enable Hanse single-position sampling. + // Probability to accept a chunk is + // min(1.0, num_records / hanse_sampling_threshold) ** hanse_sampling_gamma + // where num_records is the number of frames in the chunk. + optional uint64 hanse_sampling_threshold = 6; // by default, disabled. + optional double hanse_sampling_gamma = 7 [default = 1.0]; + optional PositionSamplingConfig position_sampling = 11; + // Optional output queue for cache hit frames. + optional QueueConfig cachehit_output = 8; + // Size of the position cache per chunk. + optional uint64 position_cache_size = 9; + // Threads for caching positions. + optional uint64 caching_threads = 10 [default = 1]; +} + +// Configuration for chunk rescorer that adjusts chunk metadata using +// tablebases. Maps to ChunkRescorerOptions in +// csrc/loader/stages/chunk_rescorer.h +message ChunkRescorerConfig { + // Number of worker threads for rescoring. + optional uint64 threads = 1 [default = 1]; + // Output queue configuration. + optional QueueConfig output = 2; + // Comma-separated list of Syzygy tablebase directories. + optional string syzygy_paths = 3; + // Policy reshaping temperature. + optional float dist_temp = 4 [default = 1.0]; + // Policy offset applied during rescoring. + optional float dist_offset = 5 [default = 0.0]; + // DTZ boost factor when policy adjustments apply. + optional float dtz_boost = 6 [default = 0.0]; + // Optional conversion target for input format (-1 keeps original). + optional int32 new_input_format = 7 [default = -1]; + // Optional deblunder threshold for policy adjustments. + optional float deblunder_threshold = 8; + // Optional deblunder width controlling smoothing around the threshold. + optional float deblunder_width = 9; + // Comma-separated list of Gaviota tablebase directories. + optional string gaviota_paths = 11; + // Soft threshold for st_q when converting v6 to v7. + optional float st_q_theta = 12 [default = 0.8333333333]; +} + +// Configuration for chunk unpacker that extracts frames from packed chunks. +// Maps to ChunkUnpackerOptions in csrc/loader/chunk_feed/chunk_unpacker.h +message ChunkUnpackerConfig { + // Number of worker threads for unpacking. + optional uint64 threads = 1 [default = 1]; + // Probability of sampling each position within a chunk. + optional float position_sampling_rate = 2 [default = 1.0]; + // Number of positions to take from each chunk. + optional uint32 position_count = 3; + // Output queue configuration. + optional QueueConfig output = 4; + // Number of positions to prefetch for caching. + optional uint32 prefetch_count = 5; + // Optional output queue for prefetch cache requests. + optional QueueConfig prefetch_output = 6; + + optional PositionSamplingConfig position_sampling = 7; +} + +// Configuration for shuffling frame sampler using reservoir sampling. +// Maps to ShufflingFrameSamplerOptions in csrc/loader/shuffling_frame_sampler.h +message ShufflingFrameSamplerConfig { + // Number of worker threads. + optional uint64 threads = 1 [default = 1]; + // Size of sampling reservoir per thread. + optional uint64 reservoir_size_per_thread = 2 [default = 1000000]; + // Output queue configuration. + optional QueueConfig output = 3; +} + +// Configuration for tensor generator that converts frames to batched tensors. +// Maps to TensorGeneratorOptions in csrc/loader/tensor_generator.h +message TensorGeneratorConfig { + // Number of worker threads for tensor generation. + optional uint64 threads = 1 [default = 1]; + // Batch size for tensor generation. + optional uint64 batch_size = 2 [default = 1024]; + // Output queue configuration. + optional QueueConfig output = 3; +} + +// Configuration for a stage that splits incoming ChunkSources into several +// outputs based on a deterministic hash of (chunk source name, index within +// source). For each configured output, provides a name, weight that controls +// the proportion of indices assigned to it, and queue parameters. +message ChunkSourceSplitterConfig { + // List of output queue configurations. + repeated QueueConfig output = 1; + // Relative weights for hash-based assignment (parallel to output). + repeated uint64 weight = 2; +} + +// Configuration for simple chunk shuffler that processes one chunk source at a +// time and outputs all chunks in shuffled order. +message SimpleChunkExtractorConfig { + // Output queue configuration. + optional QueueConfig output = 1; +} + +// Configuration for join stage that merges multiple input streams. +message JoinPositionsConfig { + // Output queue configuration. + optional QueueConfig output = 1; +} + +// Stage-level configuration providing a name and stage-specific options. +message StageConfig { + // Unique name used to reference the stage output. + optional string name = 1; + // Names of upstream stages providing input to this stage. + repeated string input = 2; + // Field 3 reserved for future output configuration. + optional FilePathProviderConfig file_path_provider = 4; + optional ChunkSourceLoaderConfig chunk_source_loader = 5; + optional ShufflingChunkPoolConfig shuffling_chunk_pool = 6; + optional ChunkUnpackerConfig chunk_unpacker = 7; + optional ShufflingFrameSamplerConfig shuffling_frame_sampler = 8; + optional TensorGeneratorConfig tensor_generator = 9; + optional ChunkRescorerConfig chunk_rescorer = 10; + optional ChunkSourceSplitterConfig chunk_source_splitter = 11; + optional SimpleChunkExtractorConfig simple_chunk_extractor = 12; + optional JoinPositionsConfig join_positions = 13; +} + +// Main configuration class for the DataLoader containing all component +// configurations. +message DataLoaderConfig { + // Ordered list of stage configurations comprising the pipeline. + repeated StageConfig stage = 1; + // Exposed outputs, each with an optional alias used by clients to fetch + // data. Expected format: "alias:stage.output", where "alias:" and ".output" + // are optional. + repeated string output = 2; +} diff --git a/proto/export_config.proto b/proto/export_config.proto new file mode 100644 index 00000000..b0209bb5 --- /dev/null +++ b/proto/export_config.proto @@ -0,0 +1,15 @@ +syntax = "proto2"; + +package lczero.training; + +// Configuration for model export settings. +message ExportConfig { + // Destination filenames where exported models will be saved. + // Supports formatting with {datetime} (YYYYMMDD-HHMMSS) and {step} (int). + // Example: "/path/to/models/lc0-{datetime}-{step:08d}.pb.gz" + repeated string destination_filename = 1; + // Training run ID for uploading to training website. Only uploads when set. + optional int32 upload_training_run = 2; + // Whether to export the SWA model instead of the main model. + optional bool export_swa_model = 3; +} diff --git a/proto/metrics_config.proto b/proto/metrics_config.proto new file mode 100644 index 00000000..8587d94a --- /dev/null +++ b/proto/metrics_config.proto @@ -0,0 +1,45 @@ +syntax = "proto2"; + +package lczero.training; + +// Sentinel message for training batch sample type. +message TrainingBatch {} + +// Configuration for weight-based metrics. +message WeightsMetric { + optional bool rms = 1; +} + +// Configuration for individual metric collection. +message MetricConfig { + // Name of the metric. + optional string name = 1; + // Whether to use the SWA model for this metric. + optional bool use_swa_model = 2; + // Period for metric collection. + optional int32 period = 3 [default = 1]; + // Whether to use global steps instead of epoch-relative steps. + optional bool use_global_steps = 4 [default = false]; + // Whether to collect metric after epoch completion. + optional bool after_epoch = 5; + + // Source of data for metric calculation. + oneof sample { + // Use training batch data. + TrainingBatch training_batch = 6; + // Name of dataloader output to use. + string dataloader_output = 7; + // Path to numpy tensor file (.npz). + string npz_filename = 8; + // Weight-based metrics. + WeightsMetric weights = 9; + } +} + +// Configuration for metrics collection and export. +message MetricsConfig { + // Directory path where TensorBoard event files will be written. + optional string tensorboard_path = 1; + // List of metrics to collect during training. + repeated MetricConfig metric = 2; +} diff --git a/proto/model_config.proto b/proto/model_config.proto new file mode 100644 index 00000000..69ed2363 --- /dev/null +++ b/proto/model_config.proto @@ -0,0 +1,61 @@ +syntax = "proto2"; + +import "proto/net.proto"; +import "proto/hlo.proto"; + +package lczero.training; + +message ModelConfig { + optional DefaultsConfig defaults = 4; + optional EmbeddingConfig embedding = 5; + optional EncoderConfig encoder = 6; + optional uint32 shared_policy_embedding_size = 7; + repeated PolicyHeadConfig policy_head = 8; + repeated ValueHeadConfig value_head = 9; + repeated MovesLeftHeadConfig movesleft_head = 10; +} + +message DefaultsConfig { + optional pblczero.XlaShapeProto.Type compute_dtype = 2; + optional pblczero.NetworkFormat.ActivationFunction activation = 4; + optional pblczero.NetworkFormat.ActivationFunction ffn_activation = 5; +} + +message EmbeddingConfig { + optional uint32 dense_size = 1; + optional uint32 embedding_size = 2; + optional uint32 dff = 3; +} + +message EncoderConfig { + optional uint32 num_blocks = 3; + optional uint32 dff = 4; + optional uint32 d_model = 5; + optional uint32 heads = 6; + optional SmolgenConfig smolgen = 9; +} + +message SmolgenConfig { + optional uint32 hidden_channels = 1; + optional uint32 hidden_size = 2; + optional uint32 gen_size = 3; + optional pblczero.NetworkFormat.ActivationFunction activation = 5; +} + +message PolicyHeadConfig { + optional string name = 1; + optional uint32 embedding_size = 2; + optional uint32 d_model = 3; +} + +message ValueHeadConfig { + optional string name = 1; + optional uint32 num_channels = 2; + optional bool has_error_output = 3; + optional uint32 num_categorical_buckets = 4; +} + +message MovesLeftHeadConfig { + optional string name = 1; + optional uint32 num_channels = 2; +} diff --git a/proto/root_config.proto b/proto/root_config.proto new file mode 100644 index 00000000..3fd4c125 --- /dev/null +++ b/proto/root_config.proto @@ -0,0 +1,27 @@ +syntax = "proto2"; + +package lczero.training; + +import "proto/data_loader_config.proto"; +import "proto/model_config.proto"; +import "proto/training_config.proto"; +import "proto/metrics_config.proto"; +import "proto/export_config.proto"; + +// Root configuration message containing all subsystem configurations. +message RootConfig { + // Name identifier for this configuration + optional string name = 1; + // Optional log filename for file-based logging + optional string log_filename = 2; + // Data loader configuration + optional DataLoaderConfig data_loader = 3; + // Training configuration + optional TrainingConfig training = 4; + // Model configuration + optional ModelConfig model = 5; + // Metrics configuration + optional MetricsConfig metrics = 6; + // Export configuration + optional ExportConfig export = 7; +} \ No newline at end of file diff --git a/proto/stage_control.proto b/proto/stage_control.proto new file mode 100644 index 00000000..bedf7827 --- /dev/null +++ b/proto/stage_control.proto @@ -0,0 +1,21 @@ +syntax = "proto2"; + +package lczero.training; + +message ShufflingChunkPoolControlRequest { + optional bool reset_chunk_anchor = 1; + optional string set_chunk_anchor = 2; +} + +message StageControlRequest { + optional ShufflingChunkPoolControlRequest chunk_pool_request = 1; +} + +message ShufflingChunkPoolControlResponse { + optional string chunk_anchor = 1; + optional int32 chunks_since_anchor = 2; +} + +message StageControlResponse { + optional ShufflingChunkPoolControlResponse chunk_pool_response = 1; +} diff --git a/proto/training_config.proto b/proto/training_config.proto new file mode 100644 index 00000000..9e4c66d3 --- /dev/null +++ b/proto/training_config.proto @@ -0,0 +1,196 @@ +syntax = "proto3"; + +package lczero.training; + +// Configuration for training algorithm and parameters. +message TrainingConfig { + ScheduleConfig schedule = 1; + repeated LrSchedule lr_schedule = 2; + CheckpointConfig checkpoint = 3; + OptimizerConfig optimizer = 4; + LossConfig losses = 5; + // Maximum gradient norm; set to 0 or omit to disable clipping. + float max_grad_norm = 6; + // Stochastic Weight Averaging (SWA) configuration. When absent, SWA is + // disabled. + SWAConfig swa = 7; +} + +message ScheduleConfig { + int32 steps_per_network = 1; + int32 chunks_per_network = 2; +} + +message OptimizerConfig { + oneof optimizer_type { + NadamwOptimizerConfig nadamw = 1; + NadamOptimizerConfig nadam = 2; + SgdOptimizerConfig sgd = 3; + } + // When set, weights selected by this selector are frozen (no gradient + // updates). + WeightsSelector freeze_selector = 4; +} + +message LrSchedule { + // Optimizer step when this schedule becomes active. + int32 starting_step = 1; + // Duration of each interval while this schedule is active. Last entry may be + // zero to indicate an open-ended tail. + repeated uint32 duration_steps = 2; + // Learning rate at the beginning of each interval. + repeated float lr = 3; + enum Transition { + CONSTANT = 0; + LINEAR = 1; + COSINE = 2; + } + // Transition type to use for each interval. Missing entries default to + // CONSTANT. + repeated Transition transition = 4; + // When true this schedule loops after finishing the final interval. + bool loop = 5; +} + +message NadamwOptimizerConfig { + float beta_1 = 1; + float beta_2 = 2; + float epsilon = 3; + float weight_decay = 4; + // Selector for weight decay. True = apply decay to this category. + WeightsSelector decay_selector = 5; +} + +message NadamOptimizerConfig { + float beta_1 = 1; + float beta_2 = 2; + float epsilon = 3; +} + +message SgdOptimizerConfig { + float momentum = 1; + bool nesterov = 2; +} + +message WeightsSelectorRule { + // Glob pattern matched against weight path. + string match = 1; + // True = include; false = exclude. + bool include = 2; +} + +// Selector for which weight categories to include. +// Rules are evaluated in order; the first matching rule applies. +// If no rule matches, otherwise_include is used. +message WeightsSelector { + repeated WeightsSelectorRule rule = 1; + bool otherwise_include = 2; +} + +message CheckpointConfig { + string path = 1; + int32 max_to_keep = 2; +} + +message LossConfig { + repeated PolicyLossConfig policy = 1; + repeated ValueLossConfig value = 2; + repeated MovesLeftLossConfig movesleft = 3; + repeated ValueErrorLossConfig value_error = 4; + repeated ValueCategoricalLossConfig value_categorical = 5; + repeated RegularizationLossConfig regularization = 6; +} + +message RegularizationLossConfig { + enum RegularizationType { + L2 = 0; + } + RegularizationType type = 1; + string metric_name = 2; + float weight = 3; + // Selector for which weights to regularize. True = include in regularization. + WeightsSelector selector = 4; +} + +enum ValueType { + RESULT = 0; + BEST = 1; + PLAYED = 2; + ORIG = 3; + ROOT = 4; + ST = 5; +} + +message OptimisticPolicyConfig { + // Name of the value head to use for optimism computation. + string value_head_name = 1; + // Which value type to use from the training sample. + ValueType value_type = 2; + // Z-score threshold for optimism (e.g. 2.0). + float strength = 3; + // Epsilon for numerical stability (e.g. 1e-5). + float eps = 4; + // Sigmoid scaling factor (e.g. 3.0). + float alpha = 5; + // If false, prevent gradients from flowing to value and error heads. + bool propagate_value_gradients = 6; +} + +message PolicyLossConfig { + string head_name = 1; + string metric_name = 2; + float weight = 3; + enum IllegalMoveHandling { + TRAIN_TO_ZERO = 0; + MASK = 1; + } + IllegalMoveHandling illegal_moves = 4; + enum LossType { + LOSS_TYPE_UNSPECIFIED = 0; + CROSS_ENTROPY = 1; + KL = 2; + } + // Selects which policy loss implementation to use. Must be specified. + LossType type = 5; + // Soft policy temperature applied before normalizing targets for KL loss. + // Values <= 0 disable the adjustment and keep raw targets. + float temperature = 6; + OptimisticPolicyConfig optimistic = 7; +} + +message ValueLossConfig { + string head_name = 1; + string metric_name = 2; + float weight = 3; + ValueType value_type = 4; +} + +message MovesLeftLossConfig { + string head_name = 1; + string metric_name = 2; + float weight = 3; + ValueType value_type = 4; +} + +message ValueErrorLossConfig { + string head_name = 1; + string metric_name = 2; + float weight = 3; + ValueType value_type = 4; + bool propagate_value_gradients = 5; +} + +message ValueCategoricalLossConfig { + string head_name = 1; + string metric_name = 2; + float weight = 3; + ValueType value_type = 4; +} + +// Stochastic Weight Averaging configuration. +message SWAConfig { + // Number of steps between SWA updates. + uint32 period_steps = 1; + // Maximum effective number of model snapshots to average. + uint32 num_averages = 2; +} diff --git a/proto/training_metrics.proto b/proto/training_metrics.proto new file mode 100644 index 00000000..5339679c --- /dev/null +++ b/proto/training_metrics.proto @@ -0,0 +1,68 @@ +syntax = "proto2"; + +package lczero; + +// Load metric that accumulates seconds of load time. +// Separate proto to support LoadMetricUpdaterProto functionality. +message LoadMetricProto { + optional string name = 1; + optional double load_seconds = 2 [default = 0.0]; + optional double total_seconds = 3 [default = 0.0]; +} + +// Statistics metric for integer values. +message StatisticsProtoInt64 { + optional int64 min = 1 [default = 9223372036854775807]; + optional int64 max = 2 [default = -9223372036854775807]; + optional int64 sum = 3 [default = 0]; + optional int64 count = 4 [default = 0]; + optional int64 latest = 5 [default = 0]; +} + +// Statistics metric for double values. +message StatisticsProtoDouble { + optional string name = 1; + optional double min = 2 [default = 1.7976931348623157e+308]; + optional double max = 3 [default = -1.7976931348623157e+308]; + optional double sum = 4 [default = 0.0]; + optional int64 count = 5 [default = 0]; + optional double latest = 6 [default = 0.0]; +} + +// Metrics for queue performance monitoring. +message QueueMetricProto { + optional string name = 1; + optional uint64 put_count = 2 [default = 0]; + optional uint64 get_count = 3 [default = 0]; + optional uint64 drop_count = 4 [default = 0]; + optional StatisticsProtoInt64 queue_fullness = 5; + optional uint64 queue_capacity = 6 [default = 0]; +} + +message CountMetricProto { + optional string name = 1; + optional uint64 count = 2 [default = 0]; +} + +// Gauge metric for values that represent current state. +message GaugeMetricProto { + optional string name = 1; + optional uint64 value = 2 [default = 0]; + optional uint64 capacity = 3; +} + +// Top-level metrics for the DataLoader. +message StageMetricProto { + optional string name = 1; + repeated LoadMetricProto load_metrics = 3; + repeated QueueMetricProto queue_metrics = 4; + repeated CountMetricProto count_metrics = 5; + repeated GaugeMetricProto gauge_metrics = 11; + repeated StatisticsProtoDouble statistics_metrics = 12; + optional string last_chunk_key = 8; + optional string anchor = 9; +} + +message DataLoaderMetricsProto { + repeated StageMetricProto stage_metrics = 1; +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..3b863915 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,110 @@ +[project] +name = "lczero-training" +version = "0.1.0" +description = "Training scripts and data loading for Leela Chess Zero" +authors = [{ name = "Leela Chess Zero Team" }] +readme = "README.md" +requires-python = ">=3.13" +dependencies = [ + "pybind11>=2.10.0", + "numpy>=1.24.0", + "textual[dev]>=0.47.0", + "protobuf==6.33.5", + "mypy>=1.17.1", + "pytest>=8.4.1", + "anyio>=4.10.0", + "jax[cuda12]==0.9.1", + "flax>=0.11.1", + "optax>=0.2.5", + "orbax-checkpoint>=0.11.23", + "python-dotenv>=1.1.1", + "requests[socks]>=2.32.5", + "matplotlib>=3.10.6", + "jaxlib==0.9.1", + "onnxruntime>=1.23.1", + "onnxruntime-gpu>=1.23.0", + "tensorboardx>=2.6.4", + "graphviz>=0.21", + "meson>=1.10.1", + "ruff>=0.15.4", +] + +[project.scripts] +lc0-leela2jax = "lczero_training.commands.leela2jax:main" +lc0-jax2leela = "lczero_training.commands.jax2leela:main" +lc0-describe = "lczero_training.commands.describe_training:main" +lc0-test-dataloader = "lczero_training.commands.test_dataloader:main" +lc0-dataloader-viz = "lczero_training.commands.dataloader_viz:main" +lc0-overfit = "lczero_training.commands.overfit:main" +lc0-init = "lczero_training.commands.training_init:main" +lc0-eval = "lczero_training.commands.training_eval:main" +lc0-tune-lr = "lczero_training.commands.tune_lr:main" +lc0-migrate-checkpoint = "lczero_training.commands.migrate_checkpoint:main" +lc0-backfill-metrics = "lczero_training.commands.backfill_metrics:main" +lc0-train = "lczero_training.commands.train:main" +lc0-daemon = "lczero_training.commands.daemon:main" +lc0-tui = "lczero_training.commands.tui:main" +lc0-weights = "lczero_training.commands.weights_tool:main" + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "mypy>=1.0.0", + "typing-extensions>=4.0.0", + "mypy-extensions>=0.4.0", + "pathspec>=0.11.0", +] + +[build-system] +requires = ["setuptools>=64", "pybind11>=2.10.0", "wheel"] +build-backend = "setuptools.build_meta" + + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-data] +"*" = ["*.so", "*.dll", "*.dylib"] + +[dependency-groups] +dev = [ + "grpcio-tools>=1.76.0", + "mypy-protobuf>=3.6.0", + "textual-dev>=1.7.0", + "types-protobuf>=6.30.2.20250809", + "types-requests>=2.32.4.20250913", +] + +[tool.mypy] +mypy_path = "src" +# TEMPORARY: Disable misc errors due to protobuf bug in PR #23156 +# (https://github.com/protocolbuffers/protobuf/pull/23156) +# Revert when protobuf fixes __slots__ = () in generated .pyi files +disable_error_code = ["misc"] + +[[tool.mypy.overrides]] +module = "orbax.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "optax" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "onnxruntime" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "tensorboardX" +ignore_missing_imports = true + + +[tool.pytest.ini_options] +testpaths = ["src"] +python_files = ["test_*.py", "*_test.py"] +pythonpath = ["src"] +addopts = "-v" + +[tool.ruff] +exclude = ["*_pb2.py*"] +line-length = 80 diff --git a/scripts/split.sh b/scripts/split.sh index a3a9951a..5f0975b7 100755 --- a/scripts/split.sh +++ b/scripts/split.sh @@ -12,6 +12,7 @@ function usage() echo " -o --output The output directory" echo " -n --window window size of test + train" echo " -t --train The training percentage in {1,...,100}" + echo " -l --latest The file to store the largest training file number." echo "" echo "Example: ./split.sh -i=/tmp -o=/out -n=2000 -t=95" echo "" @@ -39,6 +40,9 @@ do -t | --train) TRAINPCT=$VALUE ;; + -l | --latest) + LATESTFILE=$VALUE + ;; *) echo "ERROR: unknown parameter \"$PARAM\"" usage @@ -71,6 +75,12 @@ max_test=$(echo "scale=1;(1 - $TRAINPCT / 100) * $max" | bc | cut -d'.' -f1) overhead_train=$(echo "scale=1;($TRAINPCT / 100) * $overhead" | bc | cut -d'.' -f1) overhead_test=$(echo "scale=1;(1 - $TRAINPCT / 100) * $overhead" | bc | cut -d'.' -f1) +latest=0 +if [ -f $LATESTFILE ] +then + latest=$(cat $LATESTFILE) +fi + echo "" echo "start splitter, found $n games, $n_test test, $n_train train" echo " max chunks: $max" @@ -101,6 +111,15 @@ process() { id=$(echo $file | cut -d'.' -f 2) let hash_index="$id % 100 + 1" + if [ $id -gt $latest ] + then + latest=$id + if [ -f $LATESTFILE ] + then + echo $latest > $LATESTFILE + fi + fi + if [ $hash_index -gt $TRAINPCT ] then let "n_test++" diff --git a/src/lczero_training/__init__.py b/src/lczero_training/__init__.py new file mode 100644 index 00000000..c21a3922 --- /dev/null +++ b/src/lczero_training/__init__.py @@ -0,0 +1 @@ +"""Leela Chess Zero training package.""" diff --git a/src/lczero_training/_lczero_training.pyi b/src/lczero_training/_lczero_training.pyi new file mode 100644 index 00000000..259c9de7 --- /dev/null +++ b/src/lczero_training/_lczero_training.pyi @@ -0,0 +1,34 @@ +# ABOUTME: Type stubs for C++ DataLoader PyBind11 bindings. +# ABOUTME: Provides type annotations for _lczero_training compiled module. + +from typing import List, Optional, Tuple + +import numpy as np + +from proto.data_loader_config_pb2 import DataLoaderConfig +from proto.stage_control_pb2 import StageControlRequest, StageControlResponse + +class TensorBase: + def shape(self) -> List[int]: ... + def strides(self) -> List[int]: ... + def element_size(self) -> int: ... + def py_format(self) -> str: ... + +class DataLoader: + def __init__(self, config: DataLoaderConfig | bytes) -> None: ... + def add_stages(self, config: DataLoaderConfig | bytes) -> None: ... + def send_control_message( + self, request: StageControlRequest | bytes + ) -> List[Tuple[str, StageControlResponse]]: ... + def start(self) -> None: ... + def get_next(self, alias: str = "") -> Tuple[np.ndarray, ...]: ... + def maybe_get_next( + self, alias: str = "" + ) -> Optional[Tuple[np.ndarray, ...]]: ... + def stop(self) -> None: ... + def get_bucket_metrics( + self, time_period: int, include_pending: bool + ) -> Tuple[bytes, float]: ... + def get_aggregate_ending_now( + self, duration_seconds: float, include_pending: bool + ) -> Tuple[bytes, float]: ... diff --git a/src/lczero_training/commands/__init__.py b/src/lczero_training/commands/__init__.py new file mode 100644 index 00000000..a5911043 --- /dev/null +++ b/src/lczero_training/commands/__init__.py @@ -0,0 +1,21 @@ +"""Command entrypoint scaffolding and shared CLI helpers. + +This package will host thin wrappers for individual tools (convert, +training, daemon, tui) as they are extracted from nested module +__main__ dispatchers in subsequent phases. + +Phase 1 provides common helpers for consistent logging and argument +handling across commands without changing behaviour. +""" + +from .common import ( + add_logging_arguments, + configure_root_logging, + parse_log_level, +) + +__all__ = [ + "configure_root_logging", + "add_logging_arguments", + "parse_log_level", +] diff --git a/src/lczero_training/commands/backfill_metrics.py b/src/lczero_training/commands/backfill_metrics.py new file mode 100644 index 00000000..6ddfc74b --- /dev/null +++ b/src/lczero_training/commands/backfill_metrics.py @@ -0,0 +1,64 @@ +import argparse +import logging +import sys + +from lczero_training.commands import configure_root_logging +from lczero_training.training.backfill_metrics import backfill_metrics + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Backfill metrics for existing checkpoints." + ) + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the RootConfig textproto config.", + ) + parser.add_argument( + "--metrics", + nargs="+", + required=True, + help="Names of metrics to backfill (must be NPZ metrics).", + ) + parser.add_argument( + "--min-step", + type=int, + help="Minimum checkpoint step (inclusive) to process.", + ) + parser.add_argument( + "--max-step", + type=int, + help="Maximum checkpoint step (inclusive) to process.", + ) + parser.add_argument( + "--migration-config", + help=( + "Path to a CheckpointMigrationConfig textproto file. " + "If provided, checkpoints will be migrated before evaluation." + ), + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + # Lazy import to avoid heavy deps unless executing the command. + + backfill_metrics( + config_path=args.config, + metric_names=args.metrics, + min_step=args.min_step, + max_step=args.max_step, + migration_config_path=args.migration_config, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/common.py b/src/lczero_training/commands/common.py new file mode 100644 index 00000000..f137ae2b --- /dev/null +++ b/src/lczero_training/commands/common.py @@ -0,0 +1,56 @@ +import argparse +import logging +import os +import sys + +_DEFAULT_FORMAT = ( + "%(levelname).1s%(asctime)s.%(msecs)03d %(name)s " + "%(filename)s:%(lineno)d] %(message)s" +) +_DEFAULT_DATEFMT = "%m%d %H:%M:%S" + + +def configure_root_logging(level: int | str = logging.INFO) -> None: + """Configure root logging with a consistent, terse format. + + - Matches existing project format used in module __main__ files. + - Respects explicit level passed as int or name (e.g. "INFO"). + - Uses stderr by default. + - Forces reconfiguration to avoid duplicate handlers during nested runs. + """ + + resolved_level = parse_log_level(level) + logging.basicConfig( + level=resolved_level, + format=_DEFAULT_FORMAT, + datefmt=_DEFAULT_DATEFMT, + stream=sys.stderr, + force=True, + ) + + +def parse_log_level(level: int | str) -> int: + """Parse log level from int or string, with sane defaults. + + Accepts numeric levels or case-insensitive names like "DEBUG". + Falls back to INFO on invalid input. + """ + if isinstance(level, int): + return level + if isinstance(level, str): + name = level.strip().upper() + return getattr(logging, name, logging.INFO) + return logging.INFO + + +def add_logging_arguments(parser: argparse.ArgumentParser) -> None: + """Add common logging CLI arguments to a parser. + + Does not enable the flags by default; commands may opt-in and then call + configure_root_logging(parse_log_level(args.log_level)) if present. + """ + parser.add_argument( + "--log-level", + default=os.environ.get("LCZERO_LOG_LEVEL", "INFO"), + help="Logging level (DEBUG, INFO, WARNING, ERROR).", + ) diff --git a/src/lczero_training/commands/daemon.py b/src/lczero_training/commands/daemon.py new file mode 100644 index 00000000..255c4cb7 --- /dev/null +++ b/src/lczero_training/commands/daemon.py @@ -0,0 +1,27 @@ +import argparse +import logging +import sys + +from lczero_training.commands import configure_root_logging +from lczero_training.daemon.daemon import TrainingDaemon + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Run the training daemon.") + parser.add_argument("--memory-profile-dir", default=None) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + daemon = TrainingDaemon(memory_profile_dir=args.memory_profile_dir) + daemon.run() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/dataloader_viz.py b/src/lczero_training/commands/dataloader_viz.py new file mode 100644 index 00000000..70f2a8ef --- /dev/null +++ b/src/lczero_training/commands/dataloader_viz.py @@ -0,0 +1,106 @@ +import argparse +import sys + +from google.protobuf import text_format +from graphviz import Digraph # type: ignore + +from lczero_training.commands import configure_root_logging +from proto.root_config_pb2 import RootConfig + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Visualize the data loader pipeline as a graph." + ) + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the training config file.", + ) + parser.add_argument( + "--output", + type=str, + required=True, + help="Output path for the visualization (.svg or .png).", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging() + + parser = _build_parser() + args = parser.parse_args(argv) + + config = RootConfig() + with open(args.config, "r") as f: + text_format.Parse(f.read(), config) + + dot = Digraph(comment="DataLoader Pipeline") + dot.attr(rankdir="TB") + dot.attr( + "node", + shape="box", + fontname="monospace", + fontsize="10", + labeljust="l", + ) + + stage_names = set() + for stage in config.data_loader.stage: + stage_names.add(stage.name) + stage_text = text_format.MessageToString(stage, as_one_line=False) + br_tag = '
' + escaped_text = ( + stage_text.replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace("\n", br_tag) + ) + label = f"<{escaped_text}>" + dot.node(stage.name, label=label, shape="box") + + for input_spec in stage.input: + parts = input_spec.split(".", 1) + source_stage = parts[0] + if len(parts) == 2: + dot.edge(source_stage, stage.name, label=parts[1]) + else: + dot.edge(source_stage, stage.name) + + for output_spec in config.data_loader.output: + parts = output_spec.split(":", 1) + if len(parts) == 2: + alias, source = parts + else: + alias = output_spec + source = output_spec + + source_parts = source.split(".", 1) + source_stage = source_parts[0] + dot.node( + f"output_{alias}", + label=f"Output: {alias}", + shape="ellipse", + style="filled", + fillcolor="lightblue", + ) + if len(source_parts) == 2: + dot.edge(source_stage, f"output_{alias}", label=source_parts[1]) + else: + dot.edge(source_stage, f"output_{alias}") + + output_format = args.output.rsplit(".", 1)[-1].lower() + if output_format not in ("svg", "png"): + output_format = "svg" + + dot.render( + outfile=args.output, format=output_format, cleanup=True, engine="dot" + ) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/describe_training.py b/src/lczero_training/commands/describe_training.py new file mode 100644 index 00000000..daddeaea --- /dev/null +++ b/src/lczero_training/commands/describe_training.py @@ -0,0 +1,53 @@ +import argparse +import logging +import sys + +from lczero_training.commands import configure_root_logging + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Describe a trained model.") + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the training config file.", + ) + parser.add_argument( + "--shapes", + action="store_true", + help="Dump model shapes.", + ) + parser.add_argument( + "--values", + action="store_true", + help="Dump model values.", + ) + parser.add_argument( + "--weight_paths", + action="store_true", + help="List all weight paths.", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + # Import on demand to avoid importing heavy deps on --help. + from lczero_training.training.describe import describe + + describe( + config_filename=args.config, + shapes=args.shapes, + values=args.values, + weight_paths=args.weight_paths, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/jax2leela.py b/src/lczero_training/commands/jax2leela.py new file mode 100644 index 00000000..1a813b48 --- /dev/null +++ b/src/lczero_training/commands/jax2leela.py @@ -0,0 +1,132 @@ +import argparse +import gzip +import logging +import os +import sys + +import orbax.checkpoint as ocp +from google.protobuf import text_format + +from lczero_training.commands import configure_root_logging +from lczero_training.convert.jax_to_leela import ( + LeelaExportOptions, + jax_to_leela, +) +from lczero_training.training.state import TrainingState +from proto.root_config_pb2 import RootConfig + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Export JAX checkpoint to Leela format." + ) + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the training config file.", + ) + parser.add_argument( + "--output", + type=str, + required=True, + help="Output path for the Leela network (.pb.gz).", + ) + parser.add_argument( + "--export-swa", + action="store_true", + help="Export SWA model instead of regular model_state.", + ) + parser.add_argument( + "--min-version", + type=str, + default="0.31", + help="Minimum lc0 version for exported network (default: 0.31).", + ) + return parser + + +def jax2leela( + config_filename: str, + output_path: str, + export_swa: bool, + min_version: str, +) -> None: + config = RootConfig() + logging.info("Reading configuration from %s", config_filename) + with open(config_filename, "r") as f: + text_format.Parse(f.read(), config) + + if not config.training.checkpoint.path: + logging.error("Checkpoint path must be set in the configuration.") + sys.exit(1) + + logging.info("Loading checkpoint from %s", config.training.checkpoint.path) + checkpoint_mgr = ocp.CheckpointManager( + config.training.checkpoint.path, + options=ocp.CheckpointManagerOptions(create=True), + ) + + empty_state = TrainingState.new_from_config( + model_config=config.model, + training_config=config.training, + ) + + restored_state = checkpoint_mgr.restore( + checkpoint_mgr.latest_step(), args=ocp.args.PyTreeRestore(empty_state) + ) + assert isinstance(restored_state, TrainingState) + logging.info( + "Restored checkpoint at step %d", restored_state.jit_state.step + ) + + if export_swa: + if restored_state.jit_state.swa_state is None: + logging.error( + "SWA export requested but SWA state is None in checkpoint." + ) + sys.exit(1) + export_state = restored_state.jit_state.swa_state + logging.info("Exporting SWA model") + else: + export_state = restored_state.jit_state.model_state + logging.info("Exporting regular model") + + options = LeelaExportOptions( + min_version=min_version, + num_heads=restored_state.num_heads, + license=None, + training_steps=restored_state.jit_state.step, + ) + + logging.info("Converting to Leela format") + net = jax_to_leela(jax_weights=export_state, export_options=options) + + logging.info("Serializing network") + network_bytes = gzip.compress(net.SerializeToString()) + + os.makedirs(os.path.dirname(output_path), exist_ok=True) + logging.info("Writing network to %s", output_path) + with open(output_path, "wb") as f: + f.write(network_bytes) + + logging.info("Export complete") + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + jax2leela( + config_filename=args.config, + output_path=args.output, + export_swa=args.export_swa, + min_version=args.min_version, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/leela2jax.py b/src/lczero_training/commands/leela2jax.py new file mode 100644 index 00000000..6947d318 --- /dev/null +++ b/src/lczero_training/commands/leela2jax.py @@ -0,0 +1,72 @@ +import argparse +import sys + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Convert Leela Zero weights to JAX format." + ) + parser.add_argument( + "input", type=str, help="Path to the input Lc0 weights file." + ) + parser.add_argument( + "--output-model-config", + type=str, + help="Output path to the ModelConfig textproto.", + ) + parser.add_argument( + "--weights-dtype", + default="F32", + type=str, + help="The data type of the weights.", + ) + parser.add_argument( + "--compute-dtype", + default="F32", + type=str, + help="The data type for computation.", + ) + parser.add_argument( + "--print-model-config", + action="store_true", + help="Print the ModelConfig textproto to stdout.", + ) + parser.add_argument( + "--output-serialized-jax", + type=str, + help="Path to save the output JAX serialized state.", + ) + parser.add_argument( + "--output-leela-verification", + type=str, + help=( + "Path to save the round-trip converted Leela network (.pb.gz) for " + "verification." + ), + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + + # Import on demand to avoid importing heavy deps on --help. + from lczero_training.convert.leela_to_jax import ( + leela_to_jax_files, + ) + + leela_to_jax_files( + input_path=args.input, + weights_dtype=args.weights_dtype, + compute_dtype=args.compute_dtype, + output_modelconfig=args.output_model_config, + output_serialized_jax=args.output_serialized_jax, + output_leela_verification=args.output_leela_verification, + print_modelconfig=args.print_model_config, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/migrate_checkpoint.py b/src/lczero_training/commands/migrate_checkpoint.py new file mode 100644 index 00000000..511c03ec --- /dev/null +++ b/src/lczero_training/commands/migrate_checkpoint.py @@ -0,0 +1,90 @@ +import argparse +import logging +import sys + +from lczero_training.commands import configure_root_logging + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Migrate a checkpoint to a new training state." + ) + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the RootConfig textproto config.", + ) + parser.add_argument( + "--new_checkpoint", + help=( + "Path to save the new checkpoint to. If not set, the tool only " + "checks whether the migration rules fully cover the differences." + ), + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="If set, allows overwriting existing checkpoint.", + ) + parser.add_argument( + "--rules_file", + help=( + "Path to a CheckpointMigrationConfig textproto file containing " + "the migration rules." + ), + ) + parser.add_argument( + "--serialized-model", + action="store_true", + default=False, + help="Use serialized state for a model.", + ) + parser.add_argument( + "--checkpoint_step", + type=int, + help=( + "If set, use this step when loading from old checkpoint instead " + "of the latest." + ), + ) + parser.add_argument( + "--new_checkpoint_step", + type=int, + help=( + "If set, use this step when saving the new checkpoint instead of " + "copying the old step." + ), + ) + parser.add_argument("--dump_source_paths", action="store_true") + parser.add_argument("--dump_destination_paths", action="store_true") + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + # Lazy import to avoid heavy deps unless executing the command. + from lczero_training.training.migrate_checkpoint import ( + migrate_checkpoint, + ) + + migrate_checkpoint( + config=args.config, + new_checkpoint=args.new_checkpoint, + overwrite=args.overwrite, + rules_file=args.rules_file, + serialized_model=args.serialized_model, + checkpoint_step=args.checkpoint_step, + new_checkpoint_step=args.new_checkpoint_step, + dump_source_paths=args.dump_source_paths, + dump_destination_paths=args.dump_destination_paths, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/overfit.py b/src/lczero_training/commands/overfit.py new file mode 100644 index 00000000..a3188b1a --- /dev/null +++ b/src/lczero_training/commands/overfit.py @@ -0,0 +1,58 @@ +import argparse +import logging +import sys + +from lczero_training.commands import configure_root_logging + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Run an overfitting test on a single batch." + ) + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the training config file.", + ) + parser.add_argument( + "--num-steps", + type=int, + required=True, + help="Number of training steps to run on the fixed batch.", + ) + parser.add_argument( + "--coin-flip", + action="store_true", + help=( + "Train on two batches: first train batch A while evaluating batch B, then vice versa." + ), + ) + parser.add_argument( + "--csv-file", + type=str, + help="Optional path to write step-by-step overfit results.", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + # Import on demand to avoid importing heavy deps on --help. + from lczero_training.training.overfit import overfit + + overfit( + config_filename=args.config, + num_steps=args.num_steps, + coin_flip=args.coin_flip, + csv_file=args.csv_file, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/test_dataloader.py b/src/lczero_training/commands/test_dataloader.py new file mode 100644 index 00000000..393c491b --- /dev/null +++ b/src/lczero_training/commands/test_dataloader.py @@ -0,0 +1,52 @@ +import argparse +import logging +import sys + +from lczero_training.commands import configure_root_logging + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Fetch batches from the data loader to measure latency and throughput." + ) + ) + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the training config file.", + ) + parser.add_argument( + "--num-batches", + type=int, + default=10, + help="Number of batches to fetch from the data loader.", + ) + parser.add_argument( + "--npz-output", + type=str, + help="Optional path to store fetched batches as an .npz archive.", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + # Import on demand to avoid importing heavy deps on --help. + from lczero_training.training.dataloader_probe import probe_dataloader + + probe_dataloader( + config_filename=args.config, + num_batches=args.num_batches, + npz_output=args.npz_output, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/train.py b/src/lczero_training/commands/train.py new file mode 100644 index 00000000..1960f80f --- /dev/null +++ b/src/lczero_training/commands/train.py @@ -0,0 +1,135 @@ +import argparse +import datetime +import gzip +import logging +import os +import sys + +import orbax.checkpoint as ocp +from flax import nnx +from google.protobuf import text_format + +from lczero_training.commands import configure_root_logging +from lczero_training.convert.jax_to_leela import ( + LeelaExportOptions, + jax_to_leela, +) +from lczero_training.dataloader import make_dataloader +from lczero_training.model.loss_function import LczeroLoss +from lczero_training.model.model import LczeroModel +from lczero_training.training.lr_schedule import make_lr_schedule +from lczero_training.training.optimizer import make_gradient_transformation +from lczero_training.training.state import TrainingState +from lczero_training.training.training import Training, from_dataloader +from proto.root_config_pb2 import RootConfig + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Start a training run.") + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the training config file.", + ) + return parser + + +def train(config_filename: str) -> None: + config = RootConfig() + logging.info("Reading configuration from proto file") + with open(config_filename, "r") as f: + text_format.Parse(f.read(), config) + + if config.training.checkpoint.path is None: + logging.error("Checkpoint path must be set in the configuration.") + sys.exit(1) + + checkpoint_mgr = ocp.CheckpointManager( + config.training.checkpoint.path, + options=ocp.CheckpointManagerOptions( + create=True, + ), + ) + + logging.info("Creating state from configuration") + empty_state = TrainingState.new_from_config( + model_config=config.model, + training_config=config.training, + ) + logging.info("Restoring checkpoint") + training_state = checkpoint_mgr.restore( + None, args=ocp.args.PyTreeRestore(empty_state) + ) + logging.info("Restored checkpoint") + + model, _ = nnx.split( + LczeroModel(config=config.model, rngs=nnx.Rngs(params=42)) + ) + + assert isinstance(training_state, TrainingState) + + jit_state = training_state.jit_state + lr_sched = make_lr_schedule(config.training.lr_schedule) + optimizer_tx = make_gradient_transformation( + config.training.optimizer, + max_grad_norm=getattr(config.training, "max_grad_norm", 0.0), + lr_schedule=lr_sched, + ) + training = Training( + optimizer_tx=optimizer_tx, + graphdef=model, + loss_fn=LczeroLoss(config=config.training.losses), + swa_config=( + config.training.swa if config.training.HasField("swa") else None + ), + ) + new_state = training.run( + jit_state, + from_dataloader(make_dataloader(config.data_loader)), + config.training.schedule.steps_per_network, + ) + + if config.export.destination_filename: + date_str = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + + logging.info("Exporting network") + + options = LeelaExportOptions( + min_version="0.28", + num_heads=training_state.num_heads, + license=None, + training_steps=new_state.step, + ) + export_state = ( + new_state.swa_state + if config.export.export_swa_model + else new_state.model_state + ) + assert isinstance(export_state, nnx.State) + net = jax_to_leela(jax_weights=export_state, export_options=options) + network_bytes = gzip.compress(net.SerializeToString()) + + for destination_template in config.export.destination_filename: + destination = destination_template.format( + datetime=date_str, step=new_state.step + ) + logging.info(f"Writing network to {destination}") + os.makedirs(os.path.dirname(destination), exist_ok=True) + with open(destination, "wb") as f: + f.write(network_bytes) + logging.info(f"Finished writing network to {destination}") + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + train(config_filename=args.config) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/training_eval.py b/src/lczero_training/commands/training_eval.py new file mode 100644 index 00000000..59141764 --- /dev/null +++ b/src/lczero_training/commands/training_eval.py @@ -0,0 +1,83 @@ +import argparse +import logging +import sys + +from lczero_training.commands import configure_root_logging + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Evaluate a trained model.") + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the training config file.", + ) + parser.add_argument( + "--num-samples", + type=int, + help="Number of samples to evaluate.", + ) + parser.add_argument( + "--batch-size", + type=int, + help="Override batch size from data loader config.", + ) + parser.add_argument( + "--dump-stdout", + action="store_true", + help="Dump input/output tensors to stdout.", + ) + parser.add_argument( + "--dump-file", + type=str, + help="Dump input/output tensors to specified file.", + ) + parser.add_argument( + "--dump-shelve", + type=str, + help="Dump input/output tensors to specified shelve database.", + ) + parser.add_argument( + "--dump-json", + type=str, + help="Dump input/output tensors to specified JSON file.", + ) + parser.add_argument( + "--onnx-model", + type=str, + help="Path to an ONNX model to compare against JAX outputs.", + ) + parser.add_argument( + "--no-softmax-jax-wdl", + action="store_true", + help="Disable softmaxing the JAX WDL head before comparison.", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + # Lazy import to keep --help responsive and avoid heavy deps unless needed. + from lczero_training.training.eval import eval as eval_fn + + eval_fn( + config_filename=args.config, + num_samples=args.num_samples, + batch_size_override=args.batch_size, + dump_to_stdout=args.dump_stdout, + dump_to_file=args.dump_file, + dump_to_shelve=args.dump_shelve, + dump_to_json=args.dump_json, + onnx_model=args.onnx_model, + softmax_jax_wdl=not args.no_softmax_jax_wdl, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/training_init.py b/src/lczero_training/commands/training_init.py new file mode 100644 index 00000000..93cb2480 --- /dev/null +++ b/src/lczero_training/commands/training_init.py @@ -0,0 +1,87 @@ +import argparse +import logging +import sys + +from lczero_training.commands import configure_root_logging + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Initialize a new training run from a config file." + ) + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the training config file.", + ) + parser.add_argument( + "--lczero_model", + type=str, + help="Path to an existing lczero model to start from.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Seed for initializing model parameters.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Skip checkpoint creation.", + ) + parser.add_argument( + "--swa_initial_nets", + type=int, + default=0, + help="Initial value for num_averages in SWA state.", + ) + parser.add_argument( + "--override_training_steps", + type=int, + help="Override training step number.", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Allow overwriting existing checkpoint.", + ) + parser.add_argument( + "--no-copy-swa", + action="store_true", + help="Don't copy model weights to SWA state.", + ) + parser.add_argument( + "--ignore-config-mismatch", + action="store_true", + help="Ignore lczero model config mismatch.", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + # Lazy import to keep --help responsive and avoid heavy deps unless needed. + from lczero_training.training.init import init + + init( + config_filename=args.config, + lczero_model=args.lczero_model, + seed=args.seed, + dry_run=args.dry_run, + swa_initial_nets=args.swa_initial_nets, + override_training_steps=args.override_training_steps, + overwrite=args.overwrite, + no_copy_swa=args.no_copy_swa, + ignore_config_mismatch=args.ignore_config_mismatch, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/tui.py b/src/lczero_training/commands/tui.py new file mode 100644 index 00000000..02bf1ab9 --- /dev/null +++ b/src/lczero_training/commands/tui.py @@ -0,0 +1,31 @@ +import argparse +import logging +import sys + +import anyio + +from lczero_training.commands import configure_root_logging +from lczero_training.tui.app import TrainingTuiApp + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Training TUI runner") + TrainingTuiApp.add_arguments(parser) + return parser + + +async def _amain(args: argparse.Namespace) -> None: + app = TrainingTuiApp(args=args) + await app.run_async() + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + parser = _build_parser() + args = parser.parse_args(argv) + anyio.run(_amain, args) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/tune_lr.py b/src/lczero_training/commands/tune_lr.py new file mode 100644 index 00000000..c6f94069 --- /dev/null +++ b/src/lczero_training/commands/tune_lr.py @@ -0,0 +1,100 @@ +import argparse +import logging +import sys + +from lczero_training.commands import configure_root_logging + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Run a learning rate tuning sweep." + ) + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to the training config file.", + ) + parser.add_argument( + "--start-lr", + type=float, + required=True, + help="Starting learning rate for the sweep.", + ) + parser.add_argument( + "--num-steps", + type=int, + required=True, + help="Number of training steps to evaluate.", + ) + parser.add_argument( + "--multiplier", + type=float, + default=1.01, + help="Multiplier applied to the learning rate after each step.", + ) + parser.add_argument( + "--warmup-steps", + type=int, + default=0, + help=( + "Optional number of warmup steps to run at a fixed learning rate before " + "the exponential sweep." + ), + ) + parser.add_argument( + "--warmup-lr", + type=float, + help=( + "Learning rate to use during warmup steps. Required when --warmup-steps > 0." + ), + ) + parser.add_argument( + "--csv-output", + type=str, + help=( + "Optional path to write CSV results. Columns: lr, train_loss[, val_loss]." + ), + ) + parser.add_argument( + "--plot-output", + type=str, + help="Optional path to save a matplotlib plot of the sweep.", + ) + parser.add_argument( + "--num-test-batches", + type=int, + default=0, + help=( + "When > 0, also compute and report validation loss on this many fixed batches " + "(averaged each step). Default 0 (training loss only)." + ), + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging(logging.INFO) + + parser = _build_parser() + args = parser.parse_args(argv) + + # Lazy import to keep --help responsive and avoid heavy deps unless needed. + from lczero_training.training.tune_lr import tune_lr + + tune_lr( + config_filename=args.config, + start_lr=args.start_lr, + num_steps=args.num_steps, + multiplier=args.multiplier, + warmup_steps=args.warmup_steps, + warmup_lr=args.warmup_lr, + csv_output=args.csv_output, + plot_output=args.plot_output, + num_test_batches=args.num_test_batches, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/commands/weights_tool.py b/src/lczero_training/commands/weights_tool.py new file mode 100644 index 00000000..5e4cecd1 --- /dev/null +++ b/src/lczero_training/commands/weights_tool.py @@ -0,0 +1,116 @@ +"""CLI command for manipulating Lc0 neural network weights.""" + +import argparse +import sys + +import numpy as np + +from lczero_training.commands import configure_root_logging + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Manipulate Lc0 neural network weights." + ) + parser.add_argument( + "--expr", + type=str, + help=( + "Python expression to execute " + "(e.g., \"output = weights('A.pb') * 0.5\")" + ), + ) + parser.add_argument( + "script", + nargs="?", + help="Path to Python script file (if --expr not provided)", + ) + parser.add_argument( + "--input", + type=str, + action="append", + help="Pre-load input as NAME=PATH (e.g., --input A=net_A.pb.gz)", + ) + parser.add_argument( + "--output", + type=str, + help='Default output path if "output" variable is set', + ) + parser.add_argument( + "--encoding", + type=str, + default="FLOAT16", + choices=["LINEAR16", "FLOAT16", "BFLOAT16"], + help="Output encoding format (default: FLOAT16)", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + configure_root_logging() + + parser = _build_parser() + args = parser.parse_args(argv) + + # Lazy import to avoid heavy dependencies on --help. + from lczero_training.tools.weights_tool import load_weights, save_weights + from proto import net_pb2 + + # Build execution environment. + env = { + "np": np, + "weights": load_weights, + "save": save_weights, + "lc0": net_pb2, + } + + # Pre-load inputs. + if args.input: + for input_spec in args.input: + if "=" not in input_spec: + print( + f"Error: Invalid input spec '{input_spec}'. " + "Expected format: NAME=PATH", + file=sys.stderr, + ) + return 1 + name, path = input_spec.split("=", 1) + env[name] = load_weights(path) + + # Determine script source: --expr, file, or stdin. + if args.expr: + script = args.expr + elif args.script: + with open(args.script) as f: + script = f.read() + else: + if sys.stdin.isatty(): + print( + "Error: No script provided. Use --expr, provide script file, " + "or pipe to stdin.", + file=sys.stderr, + ) + return 1 + script = sys.stdin.read() + + # Execute script. + try: + exec(script, env) + except Exception as e: + print(f"Error executing script: {e}", file=sys.stderr) + return 1 + + # Auto-save if 'output' variable is set. + if "output" in env and args.output: + from lczero_training.tools.weight_wrappers import NetWrapper + + output = env["output"] + if isinstance(output, NetWrapper): + save_weights(output, args.output, args.encoding) + print(f"Saved result to {args.output}") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lczero_training/convert/__init__.py b/src/lczero_training/convert/__init__.py new file mode 100644 index 00000000..b5315ac6 --- /dev/null +++ b/src/lczero_training/convert/__init__.py @@ -0,0 +1 @@ +"""Convert package for Leela Chess Zero training.""" diff --git a/src/lczero_training/convert/jax_to_leela.py b/src/lczero_training/convert/jax_to_leela.py new file mode 100644 index 00000000..96ae34ca --- /dev/null +++ b/src/lczero_training/convert/jax_to_leela.py @@ -0,0 +1,117 @@ +import dataclasses +import logging +from typing import Optional, cast + +import numpy as np +from flax import nnx + +from lczero_training.convert.leela_pytree_visitor import ( + LeelaPytreeWeightsVisitor, +) +from proto import net_pb2 + +logger = logging.getLogger(__name__) + +_EMBEDDING_PLANE_TO_SCALE = 109 +_EMBEDDING_SCALE = 99.0 + + +class JaxToLeela(LeelaPytreeWeightsVisitor): + def embedding_block( + self, nnx_dict: nnx.State, weights: net_pb2.Weights + ) -> None: + embedding_kernel = cast(nnx.Param, nnx_dict["embedding"]["kernel"]) + original_values = embedding_kernel.value + arr = np.asarray(original_values).copy() + arr[_EMBEDDING_PLANE_TO_SCALE] /= _EMBEDDING_SCALE + embedding_kernel.value = arr + try: + super().embedding_block(nnx_dict=nnx_dict, weights=weights) + finally: + embedding_kernel.value = original_values + + def tensor( + self, + param: nnx.Param, + leela: net_pb2.Weights.Layer, + ) -> None: + weights = np.asarray(param.value, dtype=np.float32).T.flatten() + min_val, max_val = np.min(weights), np.max(weights) + range_val = max_val - min_val + + # Normalize to [0, 1], handling the case where all weights are equal. + normalized = np.where( + range_val > 1e-8, (weights - min_val) / range_val, 0.5 + ) + + # Scale to uint16 and convert to bytes. + quantized = np.round(normalized * 65535.0).astype(np.uint16) + leela.params = quantized.tobytes() + leela.min_val = float(min_val) + leela.max_val = float(max_val) + + assert len(leela.params) // 2 == weights.size + + def encoder_tower( + self, nnx_dict: nnx.State, weights: net_pb2.Weights + ) -> None: + for i in range(len(nnx_dict["encoders"]["layers"])): + weights.encoder.append(weights.EncoderLayer()) + return super().encoder_tower(nnx_dict=nnx_dict, weights=weights) + + +@dataclasses.dataclass +class LeelaExportOptions: + min_version: str + num_heads: int + license: Optional[str] + training_steps: Optional[int] = None + + +def jax_to_leela( + jax_weights: nnx.State, export_options: LeelaExportOptions +) -> net_pb2.Net: + lc0_weights = net_pb2.Net() + lc0_weights.magic = 0x1C0 + if export_options.license: + lc0_weights.license = export_options.license + ( + lc0_weights.min_version.major, + lc0_weights.min_version.minor, + lc0_weights.min_version.patch, + ) = _split_version(export_options.min_version) + lc0_weights.format.CopyFrom(_make_format()) + if export_options.training_steps is not None: + lc0_weights.training_params.training_steps = ( + export_options.training_steps + ) + + visitor = JaxToLeela(jax_weights, lc0_weights) + lc0_weights.weights.headcount = export_options.num_heads + visitor.run() + + return lc0_weights + + +def _split_version(version_str: str) -> tuple[int, int, int]: + """Splits a version string like "v12.34.56" into (12, 34, 56).""" + parts = (version_str.lstrip("v").split(".") + ["0", "0"])[:3] + return cast(tuple[int, int, int], tuple(map(int, parts))) + + +def _make_format() -> net_pb2.Format: + fmt = net_pb2.Format() + fmt.weights_encoding = fmt.LINEAR16 + netfmt = fmt.network_format + netfmt.input = netfmt.INPUT_CLASSICAL_112_PLANE + netfmt.output = netfmt.OUTPUT_WDL + netfmt.network = netfmt.NETWORK_ATTENTIONBODY_WITH_MULTIHEADFORMAT + netfmt.policy = netfmt.POLICY_ATTENTION + netfmt.value = netfmt.VALUE_WDL + netfmt.moves_left = netfmt.MOVES_LEFT_V1 + netfmt.default_activation = netfmt.DEFAULT_ACTIVATION_MISH + netfmt.smolgen_activation = netfmt.ACTIVATION_SWISH + netfmt.ffn_activation = netfmt.ACTIVATION_DEFAULT + netfmt.input_embedding = netfmt.INPUT_EMBEDDING_PE_DENSE + + return fmt diff --git a/src/lczero_training/convert/leela_pytree_visitor.py b/src/lczero_training/convert/leela_pytree_visitor.py new file mode 100644 index 00000000..21ca9e60 --- /dev/null +++ b/src/lczero_training/convert/leela_pytree_visitor.py @@ -0,0 +1,190 @@ +import math +from typing import Any, Optional + +from flax import nnx + +from proto import net_pb2 + + +class LeelaPytreeWeightsVisitor: + def __init__(self, nnx_state: nnx.State, leela_net: net_pb2.Net) -> None: + self.leela_net = leela_net + self.nnx_state = nnx_state + + def run(self) -> None: + state = self.nnx_state + weights = self.leela_net.weights + self.embedding_block(state["embedding"], weights) + self.encoder_tower(state["encoders"], weights) + self.policy_heads(state, weights.policy_heads) + for head_name in ["winner", "q", "st"]: + if head_name in state["value_heads"]: + self.value_head( + state["value_heads"][head_name], + getattr(weights.value_heads, head_name), + ) + for head_name in ["main"]: + assert head_name in state["movesleft_heads"], ( + f"movesleft head {head_name} missing in state" + ) + self.movesleft_head(state["movesleft_heads"][head_name], weights) + + def embedding_block( + self, nnx_dict: nnx.State, weights: net_pb2.Weights + ) -> None: + self.matmul( + nnx_dict["preprocess"], + weights.ip_emb_preproc_w, + weights.ip_emb_preproc_b, + ) + self.matmul( + nnx_dict["embedding"], + weights.ip_emb_w, + weights.ip_emb_b, + ) + self.layernorm( + nnx_dict["norm"], + weights.ip_emb_ln_gammas, + weights.ip_emb_ln_betas, + ) + self.tensor( + nnx_dict["ma_gating"]["mult_gate"]["gate"], weights.ip_mult_gate + ) + self.tensor( + nnx_dict["ma_gating"]["add_gate"]["gate"], weights.ip_add_gate + ) + self.ffn(nnx_dict["ffn"], weights.ip_emb_ffn) + self.layernorm( + nnx_dict["out_norm"], + weights.ip_emb_ffn_ln_gammas, + weights.ip_emb_ffn_ln_betas, + ) + + def encoder_tower( + self, nnx_dict: nnx.State, weights: net_pb2.Weights + ) -> None: + # Shared layer is stored at the point of the first usage. + self.matmul( + nnx_dict["encoders"]["layers"][0]["mha"]["smolgen"][ + "weight_gen_dense" + ], + weights.smolgen_w, + None, + ) + + # assert len(nnx_dict["encoders"]["layers"]) == len(weights.encoder) + for i in range(len(nnx_dict["encoders"]["layers"])): + self.encoder_block( + nnx_dict["encoders"]["layers"][i], weights.encoder[i] + ) + + def encoder_block( + self, nnx_dict: nnx.State, weights: net_pb2.Weights.EncoderLayer + ) -> None: + self.mha(nnx_dict["mha"], weights.mha) + self.layernorm(nnx_dict["ln1"], weights.ln1_gammas, weights.ln1_betas) + self.ffn(nnx_dict["ffn"], weights.ffn) + self.layernorm(nnx_dict["ln2"], weights.ln2_gammas, weights.ln2_betas) + + def mha(self, nnx_dict: nnx.State, weights: net_pb2.Weights.MHA) -> None: + self.matmul(nnx_dict["q"], weights.q_w, weights.q_b) + self.matmul(nnx_dict["k"], weights.k_w, weights.k_b) + self.matmul(nnx_dict["v"], weights.v_w, weights.v_b) + self.smolgen(nnx_dict["smolgen"], weights.smolgen) + self.matmul(nnx_dict["output_dense"], weights.dense_w, weights.dense_b) + + def smolgen( + self, nnx_dict: nnx.State, weights: net_pb2.Weights.Smolgen + ) -> None: + self.matmul(nnx_dict["compress"], weights.compress, None) + self.matmul(nnx_dict["dense1"], weights.dense1_w, weights.dense1_b) + self.layernorm(nnx_dict["ln1"], weights.ln1_gammas, weights.ln1_betas) + self.matmul(nnx_dict["dense2"], weights.dense2_w, weights.dense2_b) + self.layernorm(nnx_dict["ln2"], weights.ln2_gammas, weights.ln2_betas) + + def layernorm( + self, + nnx_dict: nnx.State, + scales: net_pb2.Weights.Layer, + biases: net_pb2.Weights.Layer, + ) -> None: + self.tensor(nnx_dict["scale"], scales) + self.tensor(nnx_dict["bias"], biases) + + def policy_heads( + self, nnx_dict: nnx.State, weights: net_pb2.Weights.PolicyHeads + ) -> None: + if "policy_embedding_shared" in nnx_dict: + self.matmul( + nnx_dict["policy_embedding_shared"], + weights.ip_pol_w, + weights.ip_pol_b, + ) + policy_heads_dict = nnx_dict["policy_heads"] + for head_name in ["vanilla", "optimistic_st", "soft", "opponent"]: + if head_name in policy_heads_dict: + self.policy_head( + policy_heads_dict[head_name], getattr(weights, head_name) + ) + + def policy_head( + self, nnx_dict: nnx.State, weights: net_pb2.Weights.PolicyHead + ) -> None: + if "tokens" in nnx_dict: + self.matmul(nnx_dict["tokens"], weights.ip_pol_w, weights.ip_pol_b) + self.matmul(nnx_dict["q"], weights.ip2_pol_w, weights.ip2_pol_b) + self.matmul(nnx_dict["k"], weights.ip3_pol_w, weights.ip3_pol_b) + self.matmul(nnx_dict["promotion_dense"], weights.ip4_pol_w, None) + + def value_head( + self, nnx_dict: nnx.State, weights: net_pb2.Weights.ValueHead + ) -> None: + self.matmul(nnx_dict["embed"], weights.ip_val_w, weights.ip_val_b) + self.matmul(nnx_dict["dense1"], weights.ip1_val_w, weights.ip1_val_b) + self.matmul(nnx_dict["wdl"], weights.ip2_val_w, weights.ip2_val_b) + if "error" in nnx_dict: + self.matmul( + nnx_dict["error"], weights.ip_val_err_w, weights.ip_val_err_b + ) + if "categorical" in nnx_dict: + self.matmul( + nnx_dict["categorical"], + weights.ip_val_cat_w, + weights.ip_val_cat_b, + ) + + def movesleft_head( + self, nnx_dict: nnx.State, weights: net_pb2.Weights + ) -> None: + self.matmul(nnx_dict["embed"], weights.ip_mov_w, weights.ip_mov_b) + self.matmul(nnx_dict["dense1"], weights.ip1_mov_w, weights.ip1_mov_b) + self.matmul(nnx_dict["out"], weights.ip2_mov_w, weights.ip2_mov_b) + + def ffn(self, nnx_dict: nnx.State, ffn: net_pb2.Weights.FFN) -> None: + self.matmul(nnx_dict["linear1"], ffn.dense1_w, ffn.dense1_b) + self.matmul(nnx_dict["linear2"], ffn.dense2_w, ffn.dense2_b) + + def matmul( + self, + nnx_dict: nnx.State, + weights: net_pb2.Weights.Layer, + biases: Optional[net_pb2.Weights.Layer], + ) -> None: + self.tensor(nnx_dict["kernel"], weights) + if biases: + self.tensor(nnx_dict["bias"], biases) + else: + assert "bias" not in nnx_dict + + def tensor( + self, + param: Any, + leela: net_pb2.Weights.Layer, + ) -> None: + print( + param.shape, + len(leela.params) // 2, + math.prod(param.shape), + ) + assert len(leela.params) // 2 == math.prod(param.shape) + assert len(leela.params) != 0 diff --git a/src/lczero_training/convert/leela_to_jax.py b/src/lczero_training/convert/leela_to_jax.py new file mode 100644 index 00000000..c90d12d2 --- /dev/null +++ b/src/lczero_training/convert/leela_to_jax.py @@ -0,0 +1,191 @@ +import dataclasses +import gzip +import logging +import math +from typing import Optional, cast + +import jax.numpy as jnp +from flax import nnx, serialization + +from lczero_training.model.model import LczeroModel +from proto import hlo_pb2, net_pb2 + +from .jax_to_leela import LeelaExportOptions, jax_to_leela +from .leela_pytree_visitor import LeelaPytreeWeightsVisitor +from .leela_to_modelconfig import leela_to_modelconfig + +logger = logging.getLogger(__name__) + + +_EMBEDDING_PLANE_TO_SCALE = 109 +_EMBEDDING_SCALE = 99.0 + + +@dataclasses.dataclass +class LeelaImportOptions: + weights_dtype: hlo_pb2.XlaShapeProto.Type + compute_dtype: hlo_pb2.XlaShapeProto.Type + + +def fix_older_weights_file(file: net_pb2.Net) -> None: + nf = net_pb2.NetworkFormat + has_network_format = file.format.HasField("network_format") + network_format = ( + file.format.network_format.network if has_network_format else None + ) + + net = file.format.network_format + + if not has_network_format: + # Older protobufs don't have format definition. + net.input = nf.INPUT_CLASSICAL_112_PLANE + net.output = nf.OUTPUT_CLASSICAL + net.network = nf.NETWORK_CLASSICAL_WITH_HEADFORMAT + net.value = nf.VALUE_CLASSICAL + net.policy = nf.POLICY_CLASSICAL + elif network_format == nf.NETWORK_CLASSICAL: + # Populate policyFormat and valueFormat fields in old protobufs + # without these fields. + net.network = nf.NETWORK_CLASSICAL_WITH_HEADFORMAT + net.value = nf.VALUE_CLASSICAL + net.policy = nf.POLICY_CLASSICAL + elif network_format == nf.NETWORK_SE: + net.network = nf.NETWORK_SE_WITH_HEADFORMAT + net.value = nf.VALUE_CLASSICAL + net.policy = nf.POLICY_CLASSICAL + elif ( + network_format == nf.NETWORK_SE_WITH_HEADFORMAT + and len(file.weights.encoder) > 0 + ): + # Attention body network made with old protobuf. + net.network = nf.NETWORK_ATTENTIONBODY_WITH_HEADFORMAT + if file.weights.HasField("smolgen_w"): + # Need to override activation defaults for smolgen. + net.ffn_activation = nf.ACTIVATION_RELU_2 + net.smolgen_activation = nf.ACTIVATION_SWISH + elif network_format == nf.NETWORK_AB_LEGACY_WITH_MULTIHEADFORMAT: + net.network = nf.NETWORK_ATTENTIONBODY_WITH_MULTIHEADFORMAT + + if ( + file.format.network_format.network + == nf.NETWORK_ATTENTIONBODY_WITH_HEADFORMAT + ): + weights = file.weights + if weights.HasField("policy_heads") and weights.HasField("value_heads"): + logger.info( + "Weights file has multihead format, updating format flag" + ) + net.network = nf.NETWORK_ATTENTIONBODY_WITH_MULTIHEADFORMAT + net.input_embedding = nf.INPUT_EMBEDDING_PE_DENSE + if not file.format.network_format.HasField("input_embedding"): + net.input_embedding = nf.INPUT_EMBEDDING_PE_MAP + + +class LeelaToJax(LeelaPytreeWeightsVisitor): + def embedding_block( + self, nnx_dict: nnx.State, weights: net_pb2.Weights + ) -> None: + super().embedding_block(nnx_dict=nnx_dict, weights=weights) + embedding_kernel = cast(nnx.Param, nnx_dict["embedding"]["kernel"]) + values = embedding_kernel.value + scaled_values = values.at[_EMBEDDING_PLANE_TO_SCALE].set( + values[_EMBEDDING_PLANE_TO_SCALE] * _EMBEDDING_SCALE + ) + embedding_kernel.value = scaled_values + + def tensor( + self, + param: nnx.Param, + leela: net_pb2.Weights.Layer, + ) -> None: + assert len(leela.params) // 2 == math.prod(param.shape) + assert len(leela.params) != 0 + + values = jnp.frombuffer(leela.params, dtype=jnp.uint16) + values = values.astype(jnp.float32) + alpha = values / 65535.0 + values = alpha * leela.max_val + (1.0 - alpha) * leela.min_val + values = values.astype(param.dtype) + values = values.reshape(param.shape[::-1]).transpose() + param.value = values + + +def leela_to_jax( + leela_net: net_pb2.Net, import_options: LeelaImportOptions +) -> nnx.State: + config = leela_to_modelconfig( + leela_net, + import_options.weights_dtype, + import_options.compute_dtype, + ) + + model = LczeroModel(config=config, rngs=nnx.Rngs(params=42)) + state = nnx.state(model) + visitor = LeelaToJax(state, leela_net) + visitor.run() + + return state + + +def leela_to_jax_files( + input_path: str, + weights_dtype: str, + compute_dtype: str, + output_modelconfig: Optional[str], + output_serialized_jax: Optional[str], + output_leela_verification: Optional[str], + print_modelconfig: bool = False, +) -> None: + lc0_weights = net_pb2.Net() + with gzip.open(input_path, "rb") as f: + contents = f.read() + assert isinstance(contents, bytes) + lc0_weights.ParseFromString(contents) + + fix_older_weights_file(lc0_weights) + + import_options = LeelaImportOptions( + weights_dtype=getattr(hlo_pb2.XlaShapeProto, weights_dtype), + compute_dtype=getattr(hlo_pb2.XlaShapeProto, compute_dtype), + ) + + config = leela_to_modelconfig( + lc0_weights, + import_options.weights_dtype, + import_options.compute_dtype, + ) + + if print_modelconfig: + print(config) + + if output_modelconfig: + with open(output_modelconfig, "w") as f: + f.write(str(config)) + + if output_serialized_jax is None and output_leela_verification is None: + return + + state = leela_to_jax(lc0_weights, import_options) + + if output_serialized_jax: + with open(output_serialized_jax, "wb") as f: + f.write(serialization.to_bytes(state)) + + if output_leela_verification: + min_version = ( + f"v{lc0_weights.min_version.major}." + f"{lc0_weights.min_version.minor}." + f"{lc0_weights.min_version.patch}" + ) + license_str = ( + lc0_weights.license if lc0_weights.HasField("license") else None + ) + export_options = LeelaExportOptions( + min_version=min_version, + num_heads=lc0_weights.weights.headcount, + license=license_str, + training_steps=lc0_weights.training_params.training_steps, + ) + verification_net = jax_to_leela(state, export_options) + with gzip.open(output_leela_verification, "wb") as f: + f.write(verification_net.SerializeToString()) diff --git a/src/lczero_training/convert/leela_to_modelconfig.py b/src/lczero_training/convert/leela_to_modelconfig.py new file mode 100644 index 00000000..5a24f235 --- /dev/null +++ b/src/lczero_training/convert/leela_to_modelconfig.py @@ -0,0 +1,127 @@ +from proto import hlo_pb2, model_config_pb2, net_pb2 + + +def _defaultactivation_to_activation( + activation: net_pb2.NetworkFormat.DefaultActivation, +) -> net_pb2.NetworkFormat.ActivationFunction: + return { + net_pb2.NetworkFormat.DEFAULT_ACTIVATION_RELU: net_pb2.NetworkFormat.ACTIVATION_RELU, + net_pb2.NetworkFormat.DEFAULT_ACTIVATION_MISH: net_pb2.NetworkFormat.ACTIVATION_MISH, + }[activation] + + +def leela_to_modelconfig( + leela_net: net_pb2.Net, + weights_dtype: hlo_pb2.XlaShapeProto.Type, + compute_dtype: hlo_pb2.XlaShapeProto.Type, +) -> model_config_pb2.ModelConfig: + assert weights_dtype == hlo_pb2.XlaShapeProto.F32, ( + "Only float32 weights are supported." + ) + assert leela_net.format.weights_encoding == net_pb2.Format.LINEAR16 + leela_net_format = leela_net.format.network_format + model_config = model_config_pb2.ModelConfig() + + model_config.defaults.compute_dtype = compute_dtype + model_config.defaults.activation = _defaultactivation_to_activation( + leela_net_format.default_activation + ) + model_config.defaults.ffn_activation = ( + leela_net_format.ffn_activation or model_config.defaults.activation + ) + assert ( + leela_net_format.input_embedding + == net_pb2.NetworkFormat.INPUT_EMBEDDING_PE_DENSE + ), "Only dense positional embedding is supported, got {}".format( + net_pb2.NetworkFormat.InputEmbeddingFormat.Name( + leela_net_format.input_embedding + ) + ) + assert leela_net_format.policy == net_pb2.NetworkFormat.POLICY_ATTENTION, ( + "Only attention policy is supported, got {}".format( + net_pb2.NetworkFormat.PolicyFormat.Name(leela_net_format.policy) + ) + ) + assert leela_net_format.value == net_pb2.NetworkFormat.VALUE_WDL, ( + "Only WDL value is supported, got {}".format( + net_pb2.NetworkFormat.ValueFormat.Name(leela_net_format.value) + ) + ) + assert leela_net_format.moves_left == net_pb2.NetworkFormat.MOVES_LEFT_V1, ( + "Only V1 moves left format is supported, got {}".format( + net_pb2.NetworkFormat.MovesLeftFormat.Name( + leela_net_format.moves_left + ) + ) + ) + + def size(x: net_pb2.Weights.Layer) -> int: + return len(x.params) // 2 + + assert ( + leela_net_format.network + == net_pb2.NetworkFormat.NETWORK_ATTENTIONBODY_WITH_MULTIHEADFORMAT + ) + weights = leela_net.weights + model_config.embedding.dense_size = size(weights.ip_emb_preproc_b) // 64 + model_config.embedding.embedding_size = size(weights.ip_emb_b) + assert size(weights.ip_mult_gate) > 0 + assert size(weights.ip_add_gate) > 0 + model_config.embedding.dff = size(weights.ip_emb_ffn.dense1_b) + + model_config.encoder.num_blocks = len(weights.encoder) + assert model_config.encoder.num_blocks > 0 + encoder = weights.encoder[0] + model_config.encoder.d_model = size(encoder.mha.q_b) + model_config.encoder.heads = weights.headcount + model_config.encoder.dff = size(encoder.ffn.dense1_b) + + if weights.HasField("smolgen_w"): + model_config.encoder.smolgen.activation = ( + leela_net_format.smolgen_activation + or model_config.defaults.activation + ) + model_config.encoder.smolgen.hidden_channels = ( + size(encoder.mha.smolgen.compress) + // model_config.embedding.embedding_size + ) + model_config.encoder.smolgen.gen_size = ( + size(encoder.mha.smolgen.dense2_b) // weights.headcount + ) + model_config.encoder.smolgen.hidden_size = size( + encoder.mha.smolgen.dense1_b + ) + + if weights.policy_heads.HasField("ip_pol_w"): + model_config.shared_policy_embedding_size = size( + weights.policy_heads.ip_pol_b + ) + + for head_name in ["vanilla", "optimistic_st", "soft", "opponent"]: + if weights.policy_heads.HasField(head_name): + head = getattr(weights.policy_heads, head_name) + assert size(head.ip2_pol_b) > 0 + assert not head.HasField("ip_pol_w") + policy_head = model_config.policy_head.add() + policy_head.name = head_name + if not model_config.HasField("shared_policy_embedding_size"): + policy_head.embedding_size = size(head.ip_pol_b) + policy_head.d_model = size(head.ip2_pol_b) + + for head_name in ["winner", "q", "st"]: + if weights.value_heads.HasField(head_name): + head = getattr(weights.value_heads, head_name) + assert size(head.ip_val_b) > 0 + value_head = model_config.value_head.add() + value_head.name = head_name + value_head.num_channels = size(head.ip_val_b) + if head.HasField("ip_val_err_w"): + value_head.has_error_output = True + if head.HasField("ip_val_cat_b"): + value_head.num_categorical_buckets = size(head.ip_val_cat_b) + + movesleft_head = model_config.movesleft_head.add() + movesleft_head.name = "main" + movesleft_head.num_channels = size(weights.ip_mov_b) + + return model_config diff --git a/src/lczero_training/daemon/__init__.py b/src/lczero_training/daemon/__init__.py new file mode 100644 index 00000000..9a0c8715 --- /dev/null +++ b/src/lczero_training/daemon/__init__.py @@ -0,0 +1,2 @@ +# ABOUTME: Daemon package for training subprocess communication. +# ABOUTME: Provides TrainingDaemon class for IPC via stdin/stdout. diff --git a/src/lczero_training/daemon/daemon.py b/src/lczero_training/daemon/daemon.py new file mode 100644 index 00000000..411bf815 --- /dev/null +++ b/src/lczero_training/daemon/daemon.py @@ -0,0 +1,140 @@ +import logging +import signal +import sys +import threading +import time + +import anyio + +import proto.training_metrics_pb2 as training_metrics_pb2 + +from .pipeline import TrainingPipeline +from .protocol.communicator import Communicator +from .protocol.messages import ( + StartTrainingImmediatelyPayload, + StartTrainingPayload, + TrainingStatusPayload, +) + + +class TrainingDaemon: + _training_pipeline: TrainingPipeline | None = None + _config_filepath: str | None = None + _daemon_start_time: float + + def __init__(self, memory_profile_dir: str | None = None) -> None: + self._memory_profile_dir = memory_profile_dir + self._daemon_start_time = time.time() + self._setup_logging() + self._setup_signal_handling() + self._communicator = Communicator(self, sys.stdin, sys.stdout) + self._communicator_thread = threading.Thread( + target=lambda: self._communicator.run(), daemon=True + ) + self._communicator_thread.start() + self._async_thread = threading.Thread( + target=lambda: anyio.run(self._metrics_main), daemon=True + ) + self._async_thread.start() + self._signal_thread = threading.Thread( + target=self._signal_handler_thread, daemon=True + ) + self._signal_thread.start() + + def _setup_logging(self) -> None: + logging.basicConfig( + level=logging.INFO, + format=( + "%(levelname).1s%(asctime)s.%(msecs)03d %(name)s " + "%(filename)s:%(lineno)d] %(message)s" + ), + datefmt="%m%d %H:%M:%S", + stream=sys.stderr, + ) + logging.info("TrainingDaemon starting up") + + def _setup_signal_handling(self) -> None: + # Block SIGINT and SIGTERM on all threads + signal.pthread_sigmask( + signal.SIG_BLOCK, {signal.SIGINT, signal.SIGTERM} + ) + + def _signal_handler_thread(self) -> None: + # Wait for SIGINT or SIGTERM + signum = signal.sigwait({signal.SIGINT, signal.SIGTERM}) + self._shutdown(signum) + + def _shutdown(self, signum: int) -> None: + logging.info(f"Received signal {signum}, shutting down...") + if self._training_pipeline: + self._training_pipeline.stop() + + async def _metrics_main(self) -> None: + async with anyio.create_task_group() as tg: + tg.start_soon(self._metrics_task) + + async def _metrics_task(self) -> None: + while True: + await anyio.sleep(1.1) + + dataloader_1_second = None + dataloader_total = None + dataloader_update_secs = None + training_schedule_data = None + + data_loader = None + if self._training_pipeline: + data_loader = self._training_pipeline.get_data_loader() + training_schedule_data = ( + self._training_pipeline.get_training_schedule_data( + self._daemon_start_time + ) + ) + + if data_loader is not None: + stats_1_second_bytes, _ = data_loader.get_bucket_metrics( + 0, False + ) # k1Second = 0 + stats_total_bytes, dataloader_update_secs = ( + data_loader.get_aggregate_ending_now(float("inf"), False) + ) + dataloader_1_second = ( + training_metrics_pb2.DataLoaderMetricsProto() + ) + dataloader_1_second.ParseFromString(stats_1_second_bytes) + dataloader_total = training_metrics_pb2.DataLoaderMetricsProto() + dataloader_total.ParseFromString(stats_total_bytes) + + payload = TrainingStatusPayload( + dataloader_update_secs=dataloader_update_secs, + dataloader_1_second=dataloader_1_second, + dataloader_total=dataloader_total, + training_schedule=training_schedule_data, + ) + self._communicator.send(payload) + + def run(self) -> None: + while self._config_filepath is None: + logging.info("Waiting for training config...") + time.sleep(1) + + logging.info("Config received. Starting training pipeline.") + self._training_pipeline = TrainingPipeline( + self._config_filepath, + memory_profile_dir=self._memory_profile_dir, + ) + self._training_pipeline.run() + + def on_start_training(self, payload: StartTrainingPayload) -> None: + self._config_filepath = payload.config_filepath + + def on_start_training_immediately( + self, payload: StartTrainingImmediatelyPayload + ) -> None: + if not self._training_pipeline: + logging.warning( + "Received immediate training request before pipeline initialization." + ) + return + + self._training_pipeline.start_training_immediately() diff --git a/src/lczero_training/daemon/metrics.py b/src/lczero_training/daemon/metrics.py new file mode 100644 index 00000000..4237f28f --- /dev/null +++ b/src/lczero_training/daemon/metrics.py @@ -0,0 +1,307 @@ +"""Metrics collection and logging for training daemon.""" + +import logging +import os +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Callable, Dict, Optional + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx + +from lczero_training._lczero_training import DataLoader +from lczero_training.daemon.metrics_base import _Metric +from lczero_training.daemon.rms_metrics import _RmsMetric +from lczero_training.model.loss_function import LczeroLoss +from lczero_training.training.state import JitTrainingState, TrainingSample +from lczero_training.training.tensorboard import TensorboardLogger +from lczero_training.training.training import StepHookData +from proto.metrics_config_pb2 import MetricConfig, MetricsConfig + +logger = logging.getLogger(__name__) + +# Type alias for batch tuple returned by DataLoader +# Using ... to allow variable length for compatibility with maybe_get_next +BatchTuple = tuple[np.ndarray, ...] + + +@dataclass +class CachedBatch: + """Cached batch data with the global step when it was last updated.""" + + batch: BatchTuple + global_step: int + + +def load_batch_from_npz(npz_filename: str) -> BatchTuple: + """Load a batch from an NPZ file. + + Args: + npz_filename: Path to the NPZ file. + + Returns: + BatchTuple (tuple of inputs, probabilities, values arrays). + + Raises: + ValueError: If the NPZ file doesn't contain exactly one batch. + """ + with np.load(npz_filename, allow_pickle=True) as npz_file: + batches = npz_file["batches"] + if batches.size != 1: + raise ValueError( + f"Expected 1 batch in npz '{npz_filename}', got {batches.size}" + ) + return batches[0] + + +class _TrainingBatchMetric(_Metric): + """Metric that logs training batch data.""" + + def __init__(self, config: MetricConfig, logger: TensorboardLogger): + super().__init__(config, logger) + if config.use_swa_model: + raise ValueError( + f"Metric '{config.name}': Cannot use SWA model for " + "training_batch metrics" + ) + + def log(self, hook_data: StepHookData, graphdef: nnx.GraphDef) -> None: + self.logger.log(hook_data.global_step, hook_data.metrics) + + +class _EvaluatingMetric(_Metric, ABC): + """Base class for metrics that evaluate loss on data.""" + + def __init__( + self, + config: MetricConfig, + logger: TensorboardLogger, + loss_fn: Optional[LczeroLoss], + ): + super().__init__(config, logger) + if not loss_fn: + raise ValueError(f"Metric '{config.name}': Loss function required") + self.loss_fn = loss_fn + + @abstractmethod + def get_batch(self) -> BatchTuple: + """Get the batch data to evaluate.""" + + def log(self, hook_data: StepHookData, graphdef: nnx.GraphDef) -> None: + batch = self.get_batch() + metrics = self._evaluate(batch, hook_data.jit_state, graphdef) + self.logger.log(hook_data.global_step, metrics) + + def _evaluate( + self, + batch: BatchTuple, + jit_state: JitTrainingState, + graphdef: nnx.GraphDef, + ) -> Dict[str, jax.Array]: + model_state = ( + jit_state.swa_state + if self.config.use_swa_model + else jit_state.model_state + ) + if model_state is None: + raise RuntimeError("SWA state not available") + batch_sample = TrainingSample( + inputs=jnp.asarray(batch[0]), + probabilities=jnp.asarray(batch[1]), + values=jnp.asarray(batch[2]), + ) + return _make_eval_jit(graphdef, self.loss_fn)(model_state, batch_sample) + + +_EvalJit = Callable[[nnx.State, TrainingSample], Dict[str, jax.Array]] +_eval_jit_cache: dict[tuple[int, int], _EvalJit] = {} + + +def _make_eval_jit(graphdef: nnx.GraphDef, loss_fn: LczeroLoss) -> _EvalJit: + key = (id(graphdef), id(loss_fn)) + if key not in _eval_jit_cache: + + @jax.jit + def _eval( + model_state: nnx.State, batch_sample: TrainingSample + ) -> Dict[str, jax.Array]: + model = nnx.merge(graphdef, model_state) + loss_vfn = jax.vmap(loss_fn, in_axes=(None, 0), out_axes=0) + per_sample_loss, unweighted = loss_vfn(model, batch_sample) + return { + "loss": jnp.mean(per_sample_loss), + "unweighted_losses": jax.tree_util.tree_map( + jnp.mean, unweighted + ), + } + + _eval_jit_cache[key] = _eval + return _eval_jit_cache[key] + + +def evaluate_batch( + batch: BatchTuple, + jit_state: JitTrainingState, + graphdef: nnx.GraphDef, + loss_fn: LczeroLoss, + use_swa_model: bool = False, +) -> Dict[str, jax.Array]: + """Evaluate loss function on a batch of data. + + Args: + batch: BatchTuple (inputs, probabilities, values). + jit_state: JIT training state containing model and optimizer state. + graphdef: Graph definition of the model. + loss_fn: Loss function to evaluate. + use_swa_model: If True, use SWA model state instead of regular model. + + Returns: + Dictionary of metrics with loss and unweighted losses. + """ + model_state = ( + jit_state.swa_state if use_swa_model else jit_state.model_state + ) + if model_state is None: + raise RuntimeError("SWA state not available") + batch_sample = TrainingSample( + inputs=jnp.asarray(batch[0]), + probabilities=jnp.asarray(batch[1]), + values=jnp.asarray(batch[2]), + ) + return _make_eval_jit(graphdef, loss_fn)(model_state, batch_sample) + + +class _DataLoaderMetric(_EvaluatingMetric): + """Metric that evaluates loss on dataloader output.""" + + def __init__( + self, + config: MetricConfig, + logger: TensorboardLogger, + loss_fn: Optional[LczeroLoss], + data_loader: Optional[DataLoader], + dataloader_name: str, + cached_batches: Dict[str, CachedBatch], + ): + super().__init__(config, logger, loss_fn) + if not data_loader: + raise ValueError(f"Metric '{config.name}': DataLoader required") + self.data_loader = data_loader + self.dataloader_name = dataloader_name + self.cached_batches = cached_batches + + def get_batch(self) -> BatchTuple: + return self.cached_batches[self.dataloader_name].batch + + def log(self, hook_data: StepHookData, graphdef: nnx.GraphDef) -> None: + """Update cache if needed and log the metric.""" + cached = self.cached_batches.get(self.dataloader_name) + if cached is None or cached.global_step != hook_data.global_step: + batch = self.data_loader.maybe_get_next(self.dataloader_name) + if batch is not None: + self.cached_batches[self.dataloader_name] = CachedBatch( + batch, hook_data.global_step + ) + elif cached is None: + raise RuntimeError( + f"No data for metric '{self.config.name}' " + f"from dataloader '{self.dataloader_name}'" + ) + + super().log(hook_data, graphdef) + + +class _NpzMetric(_EvaluatingMetric): + """Metric that evaluates loss on pre-loaded NPZ data.""" + + def __init__( + self, + config: MetricConfig, + logger: TensorboardLogger, + loss_fn: Optional[LczeroLoss], + npz_filename: str, + cached_batches: Dict[str, CachedBatch], + ): + super().__init__(config, logger, loss_fn) + self.npz_filename = npz_filename + self.cached_batches = cached_batches + + def get_batch(self) -> BatchTuple: + return self.cached_batches[self.npz_filename].batch + + def log(self, hook_data: StepHookData, graphdef: nnx.GraphDef) -> None: + """Load NPZ data if needed and log the metric.""" + if self.npz_filename not in self.cached_batches: + batch = load_batch_from_npz(self.npz_filename) + self.cached_batches[self.npz_filename] = CachedBatch( + batch, hook_data.global_step + ) + + super().log(hook_data, graphdef) + + +class Metrics: + """Manages metrics collection and logging for training.""" + + def __init__( + self, + config: MetricsConfig, + loss_fn: Optional[LczeroLoss] = None, + data_loader: Optional[DataLoader] = None, + ): + self._metrics: Dict[str, _Metric] = {} + self._cached_batches: Dict[str, CachedBatch] = {} + + for mc in config.metric: + tb_logger = TensorboardLogger( + os.path.join(config.tensorboard_path, mc.name) + ) + + metric: _Metric + if mc.HasField("training_batch"): + metric = _TrainingBatchMetric(mc, tb_logger) + elif mc.HasField("dataloader_output"): + metric = _DataLoaderMetric( + mc, + tb_logger, + loss_fn, + data_loader, + mc.dataloader_output, + self._cached_batches, + ) + elif mc.HasField("npz_filename"): + metric = _NpzMetric( + mc, + tb_logger, + loss_fn, + mc.npz_filename, + self._cached_batches, + ) + elif mc.HasField("weights"): + if mc.weights.rms: + metric = _RmsMetric(mc, tb_logger) + else: + raise ValueError( + f"Metric '{mc.name}': No weight metric type specified" + ) + else: + raise ValueError(f"Metric '{mc.name}' has no sample source") + + self._metrics[mc.name] = metric + + def on_step(self, hook_data: StepHookData, graphdef: nnx.GraphDef) -> None: + """Process metrics for the current step.""" + for metric in self._metrics.values(): + if metric.should_log( + hook_data.global_step, + hook_data.local_step, + hook_data.steps_per_epoch, + ): + metric.log(hook_data, graphdef) + + def close(self) -> None: + """Close all TensorBoard loggers.""" + for metric in self._metrics.values(): + metric.logger.close() diff --git a/src/lczero_training/daemon/metrics_base.py b/src/lczero_training/daemon/metrics_base.py new file mode 100644 index 00000000..ea40db8a --- /dev/null +++ b/src/lczero_training/daemon/metrics_base.py @@ -0,0 +1,30 @@ +"""Base classes for metrics.""" + +from abc import ABC, abstractmethod + +from flax import nnx + +from lczero_training.training.tensorboard import TensorboardLogger +from lczero_training.training.training import StepHookData +from proto.metrics_config_pb2 import MetricConfig + + +class _Metric(ABC): + """Base class for individual metric tracking.""" + + def __init__(self, config: MetricConfig, logger: TensorboardLogger): + self.config = config + self.logger = logger + + def should_log( + self, global_step: int, local_step: int, steps_per_epoch: int + ) -> bool: + """Check if it's time to log this metric.""" + if self.config.after_epoch and local_step + 1 == steps_per_epoch: + return True + step = global_step if self.config.use_global_steps else local_step + return (step + 1) % self.config.period == 0 + + @abstractmethod + def log(self, hook_data: StepHookData, graphdef: nnx.GraphDef) -> None: + """Log the metric for the current step.""" diff --git a/src/lczero_training/daemon/pipeline.py b/src/lczero_training/daemon/pipeline.py new file mode 100644 index 00000000..24d19055 --- /dev/null +++ b/src/lczero_training/daemon/pipeline.py @@ -0,0 +1,471 @@ +import dataclasses +import datetime +import gzip +import logging +import os +import threading +import time +from pathlib import Path +from typing import cast + +import jax +import orbax.checkpoint as ocp +import requests +from dotenv import load_dotenv +from flax import nnx +from google.protobuf import text_format + +from lczero_training._lczero_training import DataLoader +from lczero_training.convert.jax_to_leela import ( + LeelaExportOptions, + jax_to_leela, +) +from lczero_training.model.loss_function import LczeroLoss +from lczero_training.model.model import LczeroModel +from lczero_training.training.lr_schedule import make_lr_schedule +from lczero_training.training.optimizer import make_gradient_transformation +from lczero_training.training.state import JitTrainingState, TrainingState +from lczero_training.training.training import ( + StepHookData, + Training, + from_dataloader, +) +from proto.data_loader_config_pb2 import DataLoaderConfig +from proto.root_config_pb2 import RootConfig +from proto.stage_control_pb2 import StageControlRequest, StageControlResponse +from proto.training_config_pb2 import ScheduleConfig + +from .metrics import Metrics +from .protocol.messages import TrainingScheduleData, TrainingStage + +logger = logging.getLogger(__name__) + + +def _read_config_file(config_filepath: str) -> RootConfig: + config_path = Path(config_filepath) + config_text = config_path.read_text() + + root_config = RootConfig() + text_format.Parse(config_text, root_config) + return root_config + + +def _make_dataloader(config: DataLoaderConfig) -> DataLoader: + config_bytes = config.SerializeToString() + return DataLoader(config_bytes) + + +def _configure_file_logging(config: RootConfig) -> None: + """Configure file logging if log_filename is specified in config.""" + if config.HasField("log_filename"): + file_handler = logging.FileHandler(config.log_filename) + file_handler.setFormatter( + logging.Formatter( + "%(levelname).1s%(asctime)s.%(msecs)03d %(name)s " + "%(filename)s:%(lineno)d] %(message)s", + datefmt="%m%d %H:%M:%S", + ) + ) + logging.getLogger().addHandler(file_handler) + logger.info(f"Added file logging to {config.log_filename}") + + +def _log_jax_system_info() -> None: + """Log JAX system information including devices and backend details.""" + devices = jax.devices() + local_devices = jax.local_devices() + device_counts: dict[str, int] = {} + for device in devices: + device_type = device.device_kind + device_counts[device_type] = device_counts.get(device_type, 0) + 1 + + logger.info(f"JAX Backend: {jax.default_backend()}") + logger.info( + f"JAX Devices: {len(devices)} total, {len(local_devices)} local" + ) + for device_type, count in device_counts.items(): + logger.info(f" {device_type}: {count}") + for i, device in enumerate(local_devices): + logger.info(f" Local device {i}: {device}") + + +@dataclasses.dataclass +class _TrainingCycleState: + start_time: float = dataclasses.field(default_factory=time.time) + current_stage: TrainingStage = TrainingStage.WAITING_FOR_DATA + completed_epochs: int = 0 + current_cycle_start_time: float = dataclasses.field( + default_factory=time.time + ) + current_training_start_time: float | None = None + previous_training_duration: float = 0.0 + previous_cycle_duration: float = 0.0 + chunks_at_training_start: int = 0 + + +class TrainingPipeline: + _data_loader: DataLoader + _schedule: ScheduleConfig + _chunks_to_wait: int + _model: LczeroModel + _checkpoint_mgr: ocp.CheckpointManager + _training_state: TrainingState + _cycle_state: _TrainingCycleState + _metrics: Metrics | None + + def __init__( + self, + config_filepath: str, + memory_profile_dir: str | None = None, + ) -> None: + self._memory_profile_dir = memory_profile_dir + logger.info(f"Loading config from {config_filepath}") + self._config = self._load_config(config_filepath) + _configure_file_logging(self._config) + self._schedule = self._config.training.schedule + self._chunks_per_network = self._schedule.chunks_per_network + self._num_steps_per_epoch = self._schedule.steps_per_network + self._chunks_to_wait = self._chunks_per_network + self._cycle_state = _TrainingCycleState() + self._force_training_event = threading.Event() + self._metrics = None + logger.info("Creating empty model") + self._model = LczeroModel(self._config.model, rngs=nnx.Rngs(params=42)) + self._graphdef = nnx.graphdef(self._model) + logger.info( + f"Creating checkpoint manager at {self._config.training.checkpoint.path}" + ) + self._checkpoint_mgr = ocp.CheckpointManager( + self._config.training.checkpoint.path, + options=ocp.CheckpointManagerOptions( + max_to_keep=self._config.training.checkpoint.max_to_keep + or None, + ), + ) + + logger.info("Restoring checkpoint") + optimizer_config = self._config.training.optimizer + max_grad_norm = getattr(self._config.training, "max_grad_norm", 0.0) + self._lr_schedule = make_lr_schedule(self._config.training.lr_schedule) + optimizer_tx = make_gradient_transformation( + optimizer_config, + max_grad_norm=max_grad_norm, + lr_schedule=self._lr_schedule, + ) + model_state = nnx.state(self._model) + jit_state = JitTrainingState( + step=0, + model_state=model_state, + opt_state=optimizer_tx.init(model_state), + swa_state=model_state, + num_averages=0.0, + ) + empty_state = TrainingState( + jit_state=jit_state, + num_heads=self._config.model.encoder.heads, + ) + self._training_state = cast( + TrainingState, + self._checkpoint_mgr.restore( + step=None, + args=ocp.args.PyTreeRestore( + item=empty_state, + ), + ), + ) + + logger.info("Creating training session") + loss_fn = LczeroLoss(config=self._config.training.losses) + self._training = Training( + optimizer_tx=make_gradient_transformation( + self._config.training.optimizer, + max_grad_norm=max_grad_norm, + lr_schedule=self._lr_schedule, + ), + graphdef=nnx.graphdef(self._model), + loss_fn=loss_fn, + swa_config=( + self._config.training.swa + if self._config.training.HasField("swa") + else None + ), + ) + + logger.info("Creating data loader") + self._data_loader = _make_dataloader(self._config.data_loader) + self._set_chunk_anchor(self._training_state.last_chunk_source) + + # Create metrics if configured. + if self._config.HasField("metrics"): + logger.info("Creating metrics") + self._metrics = Metrics( + config=self._config.metrics, + loss_fn=loss_fn, + data_loader=self._data_loader, + ) + else: + logger.info("No metrics configured") + + _log_jax_system_info() + + def start_training_immediately(self) -> None: + """Request the next training cycle to start without waiting for chunks.""" + + logger.info("Received request to start training immediately.") + self._force_training_event.set() + + def run(self) -> None: + logging.info("Starting DataLoader") + self._data_loader.start() + + while True: + self._wait_for_chunks() + new_anchor, used_chunks = self._reset_chunk_anchor() + logging.info(f"{new_anchor=} {used_chunks=}") + self._training_state = self._training_state.replace( + last_chunk_source=new_anchor + ) + self._chunks_to_wait = max( + self._chunks_to_wait + self._chunks_per_network - used_chunks, + self._chunks_per_network // 2, + ) + self._train_one_network() + self._save_checkpoint() + network_bytes = self._export_network() + if network_bytes: + self._save_network(network_bytes) + self._upload_network(network_bytes) + + def _export_network(self) -> bytes | None: + if ( + not self._config.export.destination_filename + and not self._config.export.HasField("upload_training_run") + ): + return None + + logging.info("Exporting network") + + options = LeelaExportOptions( + min_version="0.31", + num_heads=self._training_state.num_heads, + license=None, + training_steps=self._training_state.jit_state.step, + ) + export_state = ( + self._training_state.jit_state.swa_state + if self._config.export.export_swa_model + else self._training_state.jit_state.model_state + ) + assert isinstance(export_state, nnx.State) + net = jax_to_leela(jax_weights=export_state, export_options=options) + return gzip.compress(net.SerializeToString()) + + def _save_network(self, network_bytes: bytes) -> None: + date_str = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + step = self._training_state.jit_state.step + + for destination_template in self._config.export.destination_filename: + destination = destination_template.format( + datetime=date_str, step=step + ) + logging.info(f"Writing network to {destination}") + os.makedirs(os.path.dirname(destination), exist_ok=True) + with open(destination, "wb") as f: + f.write(network_bytes) + logging.info(f"Finished writing network to {destination}") + + def _upload_network(self, network_bytes: bytes) -> None: + if not self._config.export.HasField("upload_training_run"): + return + + load_dotenv() + upload_pwd = os.getenv("UPLOAD_PWD") + if not upload_pwd: + logging.error( + "UPLOAD_PWD not found in environment variables, skipping upload." + ) + return + + try: + state = cast(nnx.State, nnx.state(self._model)) + layers = len(state["encoders"]["encoders"]["layers"]) + filters = state["embedding"]["embedding"]["bias"].shape[0] + training_id = self._config.export.upload_training_run + + logging.info( + f"Uploading network to training website (ID: {training_id}, " + f"layers: {layers}, filters: {filters})" + ) + + data = { + "pwd": upload_pwd, + "training_id": training_id, + "layers": layers, + "filters": filters, + } + response = requests.post( + "http://api.lczero.org/upload_network", + files={"file": network_bytes}, + data=data, + ) + response.raise_for_status() + + logging.info(f"Successfully uploaded network: {response.text}") + except requests.exceptions.RequestException as e: + logging.error(f"Failed to upload network: {e}") + except (KeyError, AttributeError, IndexError) as e: + logging.error(f"Failed to extract model metadata for upload: {e}") + + def _step_hook(self, hook_data: StepHookData) -> None: + # Append current learning rate from schedule to metrics. + hook_data.metrics["lr"] = self._lr_schedule(hook_data.global_step) + if self._metrics is not None: + self._metrics.on_step(hook_data, self._graphdef) + + def _train_one_network(self) -> None: + logging.info("Training one network!") + + # Record training start + self._cycle_state.current_training_start_time = time.time() + self._cycle_state.current_stage = TrainingStage.TRAINING + self._cycle_state.chunks_at_training_start = self._chunks_since_anchor() + + new_jit_state = self._training.run( + jit_state=self._training_state.jit_state, + datagen=from_dataloader(self._data_loader), + num_steps=self._schedule.steps_per_network, + step_hook=self._step_hook, + memory_profile_dir=self._memory_profile_dir, + ) + self._training_state = self._training_state.replace( + jit_state=new_jit_state + ) + + # Record training end + current_time = time.time() + if self._cycle_state.current_training_start_time: + self._cycle_state.previous_training_duration = ( + current_time - self._cycle_state.current_training_start_time + ) + self._cycle_state.previous_cycle_duration = ( + current_time - self._cycle_state.current_cycle_start_time + ) + self._cycle_state.completed_epochs += 1 + self._cycle_state.current_training_start_time = None + self._cycle_state.current_stage = TrainingStage.WAITING_FOR_DATA + self._cycle_state.current_cycle_start_time = current_time + + logging.info("Done training") + + def _save_checkpoint(self) -> None: + logging.info("Saving checkpoint") + self._checkpoint_mgr.save( + step=self._training_state.jit_state.step, + args=ocp.args.PyTreeSave(item=self._training_state), + ) + logging.info("Checkpoint saved") + + def stop(self) -> None: + self._data_loader.stop() + if self._metrics is not None: + self._metrics.close() + + def get_data_loader(self) -> DataLoader: + return self._data_loader + + def _wait_for_chunks(self) -> None: + current_chunks = self._chunks_since_anchor() + logger.info( + f"Waiting for {self._chunks_to_wait} chunks. " + f"got {current_chunks} so far" + ) + while True: + if self._force_training_event.is_set(): + logger.info( + "Force start requested; skipping remaining chunk wait." + ) + self._force_training_event.clear() + self._chunks_to_wait = self._chunks_since_anchor() + return + + if self._chunks_since_anchor() >= self._chunks_to_wait: + logger.info("Done waiting for enough chunks") + return + + time.sleep(1) + + def get_training_schedule_data( + self, daemon_start_time: float + ) -> TrainingScheduleData: + """Return current training schedule data for TUI display.""" + current_time = time.time() + + # Calculate current training time if currently training + current_training_time = 0.0 + if self._cycle_state.current_training_start_time is not None: + current_training_time = ( + current_time - self._cycle_state.current_training_start_time + ) + + # Calculate current cycle time + current_cycle_time = ( + current_time - self._cycle_state.current_cycle_start_time + ) + + # Calculate new chunks since training start + new_chunks_since_training_start = max( + 0, + self._chunks_since_anchor() + - self._cycle_state.chunks_at_training_start, + ) + + return TrainingScheduleData( + current_stage=self._cycle_state.current_stage, + completed_epochs_since_start=self._cycle_state.completed_epochs, + new_chunks_since_training_start=new_chunks_since_training_start, + chunks_to_wait=self._chunks_to_wait, + total_uptime_seconds=current_time - daemon_start_time, + current_training_time_seconds=current_training_time, + previous_training_time_seconds=self._cycle_state.previous_training_duration, + current_cycle_time_seconds=current_cycle_time, + previous_cycle_time_seconds=self._cycle_state.previous_cycle_duration, + ) + + def _send_chunk_pool_control( + self, request: StageControlRequest + ) -> StageControlResponse | None: + responses = self._data_loader.send_control_message(request) + for _, response in responses: + if response.HasField("chunk_pool_response"): + return response + return None + + def _reset_chunk_anchor(self) -> tuple[str, int]: + request = StageControlRequest() + request.chunk_pool_request.reset_chunk_anchor = True + response = self._send_chunk_pool_control(request) + if not response or not response.HasField("chunk_pool_response"): + return "", 0 + chunk_response = response.chunk_pool_response + return chunk_response.chunk_anchor, chunk_response.chunks_since_anchor + + def _chunks_since_anchor(self) -> int: + request = StageControlRequest() + request.chunk_pool_request.SetInParent() + response = self._send_chunk_pool_control(request) + if not response or not response.HasField("chunk_pool_response"): + return 0 + return response.chunk_pool_response.chunks_since_anchor + + def _set_chunk_anchor(self, anchor: str) -> None: + request = StageControlRequest() + request.chunk_pool_request.set_chunk_anchor = anchor or "" + self._send_chunk_pool_control(request) + + def _load_config(self, config_filepath: str) -> RootConfig: + config_path = Path(config_filepath) + config_text = config_path.read_text() + + root_config = RootConfig() + text_format.Parse(config_text, root_config) + return root_config diff --git a/src/lczero_training/daemon/protocol/__init__.py b/src/lczero_training/daemon/protocol/__init__.py new file mode 100644 index 00000000..6bc8ead5 --- /dev/null +++ b/src/lczero_training/daemon/protocol/__init__.py @@ -0,0 +1,2 @@ +# ABOUTME: Protocol package for JSONL IPC communication between processes. +# ABOUTME: Contains registry system, message definitions, and communicator class. diff --git a/src/lczero_training/daemon/protocol/communicator.py b/src/lczero_training/daemon/protocol/communicator.py new file mode 100644 index 00000000..95a1be0b --- /dev/null +++ b/src/lczero_training/daemon/protocol/communicator.py @@ -0,0 +1,230 @@ +# ABOUTME: Core Communicator class for JSONL IPC between processes. +# ABOUTME: Handles serialization/deserialization and message dispatch via stdin/stdout. + +import json +import types +from dataclasses import is_dataclass +from enum import Enum +from typing import Any, Optional, TextIO, Union, get_args, get_origin + +from anyio.streams.text import TextReceiveStream, TextSendStream +from google.protobuf.json_format import MessageToDict, ParseDict +from google.protobuf.message import Message + +from .registry import CLASS_TO_TYPE_MAP, TYPE_TO_CLASS_MAP + + +def _to_serializable(obj: Any) -> Any: + """Convert dataclass/protobuf objects to JSON-serializable dicts.""" + if isinstance(obj, Message): + return MessageToDict( + obj, preserving_proto_field_name=True, use_integers_for_enums=True + ) + elif isinstance(obj, Enum): + return obj.value + elif is_dataclass(obj): + return { + f.name: _to_serializable(getattr(obj, f.name)) + for f in obj.__dataclass_fields__.values() + if getattr(obj, f.name) is not None + } + elif isinstance(obj, (list, tuple)): + return [_to_serializable(item) for item in obj] + elif isinstance(obj, dict): + return {k: _to_serializable(v) for k, v in obj.items()} + else: + return obj + + +def _unwrap_optional(t: Any) -> Any: + """Extract T from T | None or Union[T, None].""" + if isinstance(t, types.UnionType) or get_origin(t) is Union: + args = [a for a in get_args(t) if a is not type(None)] + return args[0] if len(args) == 1 else t + return t + + +def _is_protobuf(cls: type) -> bool: + """Check if cls is a protobuf Message class.""" + try: + return isinstance(cls, type) and issubclass(cls, Message) + except TypeError: + return False + + +def _from_serializable(cls: type, data: Any) -> Any: + """Reconstruct dataclass/protobuf from dict.""" + if _is_protobuf(cls): + instance = cls() + ParseDict(data, instance) + return instance + + if not is_dataclass(cls): + return data + + args = {} + for field in cls.__dataclass_fields__.values(): + if field.name not in data: + continue + + value = data[field.name] + field_type = _unwrap_optional(field.type) + + if get_origin(field_type) is list: + item_type = get_args(field_type)[0] + if is_dataclass(item_type) or _is_protobuf(item_type): + value = [_from_serializable(item_type, item) for item in value] + elif is_dataclass(field_type) or _is_protobuf(field_type): + value = _from_serializable(field_type, value) + elif isinstance(field_type, type) and issubclass(field_type, Enum): + # Convert string value back to enum + value = field_type(value) + + args[field.name] = value + + return cls(**args) + + +class Communicator: + def __init__( + self, handler: Any, input_stream: TextIO, output_stream: TextIO + ) -> None: + """ + Initializes the Communicator. + + Args: + handler: An object with `on_` methods. + input_stream: A file-like object to read incoming messages from (e.g., sys.stdin). + output_stream: A file-like object to write outgoing messages to (e.g., sys.stdout). + """ + self.handler = handler + self.input = input_stream + self.output = output_stream + + def send(self, payload_instance: Any) -> None: + """ + Serializes and sends a payload object as a notification. + The event type is automatically looked up from the registry. + """ + payload_cls = type(payload_instance) + event_type = CLASS_TO_TYPE_MAP.get(payload_cls) + + if event_type is None: + raise TypeError( + f"Object of type {payload_cls.__name__} is not a registered payload." + ) + + payload_dict = _to_serializable(payload_instance) + message = {"type": event_type, "payload": payload_dict} + + self.output.write(json.dumps(message) + "\n") + self.output.flush() + + def _dispatch(self, line: str) -> None: + line = line.strip() + if not line: + return + + data = json.loads(line) + event_type = data["type"] + payload_dict = data["payload"] + + payload_cls = TYPE_TO_CLASS_MAP[event_type] + payload_instance = _from_serializable(payload_cls, payload_dict) + + handler_method_name = f"on_{event_type}" + handler_method = getattr(self.handler, handler_method_name) + + handler_method(payload_instance) + + def run(self) -> None: + """ + Starts the blocking listener loop. + Reads from the input stream line-by-line, deserializes notifications, + and dispatches them to the appropriate handler method. + + This method blocks until the input stream is closed. + """ + for line in self.input: + self._dispatch(line) + + +class AsyncCommunicator: + def __init__( + self, + handler: Any, + input_stream: TextReceiveStream, + output_stream: TextSendStream, + io_dump: Optional[TextIO] = None, + ) -> None: + """ + Initializes the AsyncCommunicator. + + Args: + handler: An object with async `on_` methods. + input_stream: A TextReceiveStream to read incoming messages from. + output_stream: A TextSendStream to write outgoing messages to. + io_dump: Optional file to dump raw IO for debugging. + """ + self.handler = handler + self.input_stream = input_stream + self.output_stream = output_stream + self._io_dump = io_dump + self._buffer = "" + + async def send(self, payload_instance: Any) -> None: + """ + Serializes and sends a payload object as a notification. + The event type is automatically looked up from the registry. + """ + payload_cls = type(payload_instance) + event_type = CLASS_TO_TYPE_MAP.get(payload_cls) + + if event_type is None: + raise TypeError( + f"Object of type {payload_cls.__name__} is not a registered payload." + ) + + payload_dict = _to_serializable(payload_instance) + message = {"type": event_type, "payload": payload_dict} + + message_line = json.dumps(message) + "\n" + if self._io_dump: + self._io_dump.write(f"> {message_line}") + await self.output_stream.send(message_line) + + async def _dispatch(self, line: str) -> None: + line = line.strip() + if not line: + return + + data = json.loads(line) + event_type = data["type"] + payload_dict = data["payload"] + + payload_cls = TYPE_TO_CLASS_MAP[event_type] + payload_instance = _from_serializable(payload_cls, payload_dict) + + handler_method_name = f"on_{event_type}" + handler_method = getattr(self.handler, handler_method_name) + + await handler_method(payload_instance) + + async def run(self) -> None: + """ + Starts the async listener loop. + Reads from the input stream line-by-line, deserializes notifications, + and dispatches them to the appropriate async handler method. + + This method runs until the input stream is closed. + """ + async for chunk in self.input_stream: + self._buffer += chunk + while "\n" in self._buffer: + line, self._buffer = self._buffer.split("\n", 1) + if self._io_dump and line: + self._io_dump.write(f"< {line}\n") + await self._dispatch(line) + + if self._buffer: + await self._dispatch(self._buffer) diff --git a/src/lczero_training/daemon/protocol/messages.py b/src/lczero_training/daemon/protocol/messages.py new file mode 100644 index 00000000..d94e2601 --- /dev/null +++ b/src/lczero_training/daemon/protocol/messages.py @@ -0,0 +1,59 @@ +# ABOUTME: Payload dataclass definitions for JSONL IPC protocol messages. +# ABOUTME: Defines minimal event types for training daemon communication. + +from dataclasses import dataclass +from enum import Enum +from typing import Optional + +import proto.training_metrics_pb2 as training_metrics_pb2 + +from .registry import register + + +class TrainingStage(Enum): + WAITING_FOR_DATA = "WAITING FOR DATA" + TRAINING = "TRAINING" + + +@dataclass +class TrainingScheduleData: + current_stage: TrainingStage + completed_epochs_since_start: int + new_chunks_since_training_start: int + chunks_to_wait: int + total_uptime_seconds: float + current_training_time_seconds: float + previous_training_time_seconds: float + current_cycle_time_seconds: float + previous_cycle_time_seconds: float + + +# --- Notifications from UI (Parent) to Trainer (Child) --- + + +@register("start_training") +@dataclass +class StartTrainingPayload: + config_filepath: str + + +@register("start_training_immediately") +@dataclass +class StartTrainingImmediatelyPayload: + pass + + +# --- Notifications from Trainer (Child) to UI (Parent) --- + + +@register("training_status") +@dataclass +class TrainingStatusPayload: + dataloader_update_secs: Optional[float] = None + dataloader_1_second: Optional[ + training_metrics_pb2.DataLoaderMetricsProto + ] = None + dataloader_total: Optional[training_metrics_pb2.DataLoaderMetricsProto] = ( + None + ) + training_schedule: Optional[TrainingScheduleData] = None diff --git a/src/lczero_training/daemon/protocol/registry.py b/src/lczero_training/daemon/protocol/registry.py new file mode 100644 index 00000000..ce22dd78 --- /dev/null +++ b/src/lczero_training/daemon/protocol/registry.py @@ -0,0 +1,33 @@ +# ABOUTME: Registry system for mapping event type strings to payload dataclasses. +# ABOUTME: Provides @register decorator and maintains bidirectional mapping dicts. + +import inspect +from typing import Callable + +# These maps will be populated by the @register decorator +TYPE_TO_CLASS_MAP = {} +CLASS_TO_TYPE_MAP = {} + + +def register(event_type: str) -> Callable[[type], type]: + """A decorator to register a payload dataclass with its event type string.""" + + def decorator(cls: type) -> type: + if not inspect.isclass(cls): + raise TypeError( + "The @register decorator can only be used on classes." + ) + + if event_type in TYPE_TO_CLASS_MAP: + raise ValueError( + f"Event type '{event_type}' is already registered." + ) + + if cls in CLASS_TO_TYPE_MAP: + raise ValueError(f"Class '{cls.__name__}' is already registered.") + + TYPE_TO_CLASS_MAP[event_type] = cls + CLASS_TO_TYPE_MAP[cls] = event_type + return cls + + return decorator diff --git a/src/lczero_training/daemon/rms_metrics.py b/src/lczero_training/daemon/rms_metrics.py new file mode 100644 index 00000000..60c74188 --- /dev/null +++ b/src/lczero_training/daemon/rms_metrics.py @@ -0,0 +1,117 @@ +"""RMS metrics for model parameters.""" + +from typing import Any, cast + +import jax +import jax.numpy as jnp +from flax import nnx + +from lczero_training.daemon.metrics_base import _Metric +from lczero_training.model.encoder import EncoderBlock +from lczero_training.model.model import LczeroModel +from lczero_training.training.tensorboard import TensorboardLogger +from lczero_training.training.training import StepHookData +from proto.metrics_config_pb2 import MetricConfig + + +@jax.jit +def compute_rms(state_subtree: nnx.State) -> jax.Array: + """Compute RMS of all parameters in a state subtree.""" + leaves = jax.tree_util.tree_leaves(state_subtree) + total_sq = sum(jnp.sum(jnp.square(p)) for p in leaves) + total_n = sum(p.size for p in leaves) + return jnp.sqrt(total_sq / total_n) + + +def extract_attention_components(model: LczeroModel) -> dict[str, Any]: + """Extract Q, K, V, output_dense, smolgen from all encoder layers. + + Args: + model: LczeroModel instance. + + Returns: + Dict with keys 'q', 'k', 'v', 'output_dense', optionally 'smolgen'. + """ + components: dict[str, Any] = { + "q": {}, + "k": {}, + "v": {}, + "output_dense": {}, + } + + encoder_layers = cast(list[EncoderBlock], model.encoders.encoders.layers) + for i, encoder_block in enumerate(encoder_layers): + mha = encoder_block.mha + components["q"][f"layer_{i}"] = nnx.state(mha.q) + components["k"][f"layer_{i}"] = nnx.state(mha.k) + components["v"][f"layer_{i}"] = nnx.state(mha.v) + components["output_dense"][f"layer_{i}"] = nnx.state(mha.output_dense) + + if mha.smolgen is not None: + if "smolgen" not in components: + components["smolgen"] = {} + components["smolgen"][f"layer_{i}"] = nnx.state(mha.smolgen) + + return components + + +def collect_rms_metrics(model: LczeroModel) -> dict[str, Any]: + """Collect all RMS metrics for the model. + + Args: + model: LczeroModel instance. + + Returns: + Nested dict with RMS values for different model components. + """ + model_state = nnx.state(model) + + metrics: dict[str, Any] = { + "all_params": compute_rms(model_state), + "embedding": compute_rms(nnx.state(model.embedding)), + "encoder_body": compute_rms(nnx.state(model.encoders)), + } + + # Attention components + attn_components = extract_attention_components(model) + metrics["attention"] = { + name: compute_rms(component) + for name, component in attn_components.items() + } + + # Policy heads + metrics["policy_heads"] = { + name: compute_rms(nnx.state(head)) + for name, head in model.policy_heads.items() + } + + # Value heads + metrics["value_heads"] = { + name: compute_rms(nnx.state(head)) + for name, head in model.value_heads.items() + } + + # Movesleft heads + metrics["movesleft_heads"] = { + name: compute_rms(nnx.state(head)) + for name, head in model.movesleft_heads.items() + } + + return metrics + + +class _RmsMetric(_Metric): + """Metric that computes RMS of model parameters.""" + + def __init__(self, config: MetricConfig, logger: TensorboardLogger): + super().__init__(config, logger) + + def log(self, hook_data: StepHookData, graphdef: nnx.GraphDef) -> None: + model_state = ( + hook_data.jit_state.swa_state + if self.config.use_swa_model + else hook_data.jit_state.model_state + ) + model = nnx.merge(graphdef, model_state) + metrics = collect_rms_metrics(model) + self.logger.log(hook_data.global_step, metrics) diff --git a/src/lczero_training/dataloader/__init__.py b/src/lczero_training/dataloader/__init__.py new file mode 100644 index 00000000..b419ce7c --- /dev/null +++ b/src/lczero_training/dataloader/__init__.py @@ -0,0 +1,13 @@ +from lczero_training._lczero_training import ( + DataLoader, + TensorBase, +) +from proto.data_loader_config_pb2 import DataLoaderConfig + +__all__ = ["DataLoader", "make_dataloader", "TensorBase"] + + +def make_dataloader(config: DataLoaderConfig) -> DataLoader: + loader = DataLoader(config) + loader.start() + return loader diff --git a/src/lczero_training/model/__init__.py b/src/lczero_training/model/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/lczero_training/model/embedding.py b/src/lczero_training/model/embedding.py new file mode 100644 index 00000000..f79dbf1f --- /dev/null +++ b/src/lczero_training/model/embedding.py @@ -0,0 +1,104 @@ +import jax +import jax.numpy as jnp +from flax import nnx + +from proto import model_config_pb2 + +from .shared import Ffn +from .utils import get_activation + + +class Embedding(nnx.Module): + """Computes embeddings for the input features.""" + + def __init__( + self, + *, + input_channels: int, + config: model_config_pb2.EmbeddingConfig, + defaults: model_config_pb2.DefaultsConfig, + deepnorm_alpha: float, + deepnorm_beta: float, + rngs: nnx.Rngs, + ): + self._input_channels = input_channels + dense_size = config.dense_size + embedding_size = config.embedding_size + self.activation = defaults.activation + + assert dense_size > 0 + self.preprocess = nnx.Linear( + in_features=64 * 12, + out_features=64 * dense_size, + rngs=rngs, + ) + + assert embedding_size > 0 + self.embedding = nnx.Linear( + in_features=input_channels + dense_size, + out_features=embedding_size, + rngs=rngs, + ) + self.norm = nnx.LayerNorm(embedding_size, epsilon=1e-3, rngs=rngs) + self.ma_gating = MaGating(feature_shape=(64, embedding_size), rngs=rngs) + self.deepnorm_alpha = deepnorm_alpha + self.ffn = Ffn( + in_features=embedding_size, + hidden_features=config.dff, + hidden_activation=defaults.ffn_activation, + deepnorm_beta=deepnorm_beta, + rngs=rngs, + ) + self.out_norm = nnx.LayerNorm(embedding_size, epsilon=1e-3, rngs=rngs) + + def __call__(self, x: jax.Array) -> jax.Array: + # Preprocess positional info and concatenate to input. + pos_info = self.preprocess(x[..., :12].flatten()).reshape((64, -1)) + x = jnp.concatenate([x, pos_info], axis=1) + + # Square embedding. + x = self.embedding(x) + x = get_activation(self.activation)(x) + x = self.norm(x) + x = self.ma_gating(x) + # FFN block with residual connection and layer norm. + x = x + self.ffn(x) * self.deepnorm_alpha + x = self.out_norm(x) + return x + + +class MaGating(nnx.Module): + """Applies multiplicative and additive gating.""" + + def __init__(self, feature_shape: tuple[int, ...], *, rngs: nnx.Rngs): + self.mult_gate = Gating( + feature_shape=feature_shape, additive=False, rngs=rngs + ) + self.add_gate = Gating( + feature_shape=feature_shape, additive=True, rngs=rngs + ) + + def __call__(self, x: jax.Array) -> jax.Array: + return self.add_gate(self.mult_gate(x)) + + +class Gating(nnx.Module): + def __init__( + self, + feature_shape: tuple[int, ...], + additive: bool = True, + *, + rngs: nnx.Rngs, + ): + self.additive = additive + init_val = 0.0 if self.additive else 1.0 + self.gate = nnx.Param( + jnp.full(feature_shape, init_val, dtype=jnp.float32) + ) + + def __call__(self, inputs: jax.Array) -> jax.Array: + if self.additive: + return inputs + self.gate.value + + effective_gate = jax.nn.relu(self.gate.value) + return inputs * effective_gate diff --git a/src/lczero_training/model/encoder.py b/src/lczero_training/model/encoder.py new file mode 100644 index 00000000..6460fb6f --- /dev/null +++ b/src/lczero_training/model/encoder.py @@ -0,0 +1,230 @@ +import math +from typing import Optional + +import jax +import jax.numpy as jnp +from flax import nnx +from flax.linen import initializers as flax_initializers + +from proto import model_config_pb2 + +from .shared import Ffn +from .utils import get_activation + + +class EncoderTower(nnx.Module): + def __init__( + self, + *, + in_features: int, + config: model_config_pb2.EncoderConfig, + defaults: model_config_pb2.DefaultsConfig, + deepnorm_beta: float, + rngs: nnx.Rngs, + ): + smolgen_shared_gen_dense = None + assert config.HasField("smolgen") + if config.HasField("smolgen"): + smolgen_shared_gen_dense = nnx.Linear( + in_features=config.smolgen.gen_size, + out_features=64 * 64, + use_bias=False, + rngs=rngs, + ) + + self.encoders = nnx.Sequential( + *[ + EncoderBlock( + in_features=in_features, + config=config, + defaults=defaults, + smol_gen_dense=smolgen_shared_gen_dense, + deepnorm_beta=deepnorm_beta, + rngs=rngs, + ) + for _ in range(config.num_blocks) + ] + ) + + def __call__(self, x: jax.Array) -> jax.Array: + return self.encoders(x) + + +class EncoderBlock(nnx.Module): + """A single block of the transformer encoder.""" + + def __init__( + self, + *, + in_features: int, + config: model_config_pb2.EncoderConfig, + defaults: model_config_pb2.DefaultsConfig, + smol_gen_dense: Optional[nnx.Linear], + deepnorm_beta: float, + rngs: nnx.Rngs, + ): + assert (smol_gen_dense is not None) == config.HasField("smolgen") + self.mha = MultiHeadAttention( + in_features=in_features, + config=config, + defaults=defaults, + smol_gen_dense=smol_gen_dense, + deepnorm_beta=deepnorm_beta, + rngs=rngs, + ) + + self.alpha = math.pow(2.0 * config.num_blocks, -0.25) + self.ln1 = nnx.LayerNorm(in_features, epsilon=1e-3, rngs=rngs) + self.ffn = Ffn( + in_features=in_features, + hidden_features=config.dff, + hidden_activation=defaults.ffn_activation, + deepnorm_beta=deepnorm_beta, + rngs=rngs, + ) + self.ln2 = nnx.LayerNorm(in_features, epsilon=1e-3, rngs=rngs) + + def __call__(self, x: jax.Array) -> jax.Array: + x = x + self.mha(x) * self.alpha + out1 = self.ln1(x) + ffn_out = self.ffn(out1) + return self.ln2(out1 + ffn_out * self.alpha) + + +class MultiHeadAttention(nnx.Module): + """Multi-head attention module.""" + + def __init__( + self, + in_features: int, + config: model_config_pb2.EncoderConfig, + defaults: model_config_pb2.DefaultsConfig, + smol_gen_dense: Optional[nnx.Linear], + deepnorm_beta: float, + *, + rngs: nnx.Rngs, + ): + depth = config.d_model + assert depth % config.heads == 0, ( + "Model depth must be divisible by the number of heads." + ) + self.activation = defaults.activation + self.depth = depth + self.num_heads = config.heads + self.q = nnx.Linear( + in_features=in_features, out_features=depth, rngs=rngs + ) + self.k = nnx.Linear( + in_features=in_features, out_features=depth, rngs=rngs + ) + deepnorm_init = flax_initializers.variance_scaling( + scale=deepnorm_beta, + mode="fan_avg", + distribution="truncated_normal", + ) + + self.v = nnx.Linear( + in_features=in_features, + out_features=depth, + kernel_init=deepnorm_init, + rngs=rngs, + ) + self.output_dense = nnx.Linear( + in_features=depth, + out_features=in_features, + kernel_init=deepnorm_init, + rngs=rngs, + ) + + assert (smol_gen_dense is not None) == config.HasField("smolgen") + self.smolgen: Optional[Smolgen] + if smol_gen_dense is not None: + self.smolgen = Smolgen( + in_features=in_features, + config=config.smolgen, + defaults=defaults, + heads=config.heads, + weight_gen_dense=smol_gen_dense, + rngs=rngs, + ) + else: + self.smolgen = None + + def __call__(self, x: jax.Array) -> jax.Array: + q, k, v = self.q(x), self.k(x), self.v(x) + + head_depth = self.depth // self.num_heads + # Reshape for multi-head attention. + q, k, v = ( + t.reshape((-1, self.num_heads, head_depth)).transpose((1, 0, 2)) + for t in (q, k, v) + ) + + # Scaled dot-product attention. + logits = jnp.einsum("...qd,...kd->...qk", q, k) + logits /= jnp.sqrt(k.shape[-1]).astype(k.dtype) + + if self.smolgen is not None: + logits += self.smolgen(x) + + attention_weights = nnx.softmax(logits, axis=-1) + scaled_attention = jnp.matmul(attention_weights, v) + + # Reshape back to original dimensions. + scaled_attention = scaled_attention.transpose((1, 0, 2)).reshape( + (-1, self.depth) + ) + return self.output_dense(scaled_attention) + + +class Smolgen(nnx.Module): + """Smolgen module for generating attention biases.""" + + def __init__( + self, + in_features: int, + config: model_config_pb2.SmolgenConfig, + defaults: model_config_pb2.DefaultsConfig, + heads: int, + weight_gen_dense: nnx.Linear, + *, + rngs: nnx.Rngs, + ): + self.heads = heads + self.compress = nnx.Linear( + in_features=in_features, + out_features=config.hidden_channels, + use_bias=False, + rngs=rngs, + ) + self.dense1 = nnx.Linear( + in_features=config.hidden_channels * 64, + out_features=config.hidden_size, + rngs=rngs, + ) + self.ln1 = nnx.LayerNorm(config.hidden_size, epsilon=1e-3, rngs=rngs) + + self.dense2 = nnx.Linear( + in_features=config.hidden_size, + out_features=config.gen_size * heads, + rngs=rngs, + ) + self.ln2 = nnx.LayerNorm( + config.gen_size * heads, epsilon=1e-3, rngs=rngs + ) + self.weight_gen_dense = weight_gen_dense + self.activation = config.activation or defaults.activation + + def __call__(self, x: jax.Array) -> jax.Array: + compressed = self.compress(x).flatten() + hidden = self.dense1(compressed) + hidden = get_activation(self.activation)(hidden) + hidden = self.ln1(hidden) + + gen_from = self.dense2(hidden) + gen_from = get_activation(self.activation)(gen_from) + gen_from = self.ln2(gen_from) + gen_from = gen_from.reshape((self.heads, -1)) + + out = self.weight_gen_dense(gen_from) + return out.reshape((self.heads, 64, 64)) diff --git a/src/lczero_training/model/loss_function.py b/src/lczero_training/model/loss_function.py new file mode 100644 index 00000000..15133143 --- /dev/null +++ b/src/lczero_training/model/loss_function.py @@ -0,0 +1,419 @@ +from typing import Dict, List, Optional, Sequence, Tuple, Union, cast + +import jax +import jax.numpy as jnp +import optax +from flax import nnx +from jax.scipy.special import xlogy + +from lczero_training.training.state import TrainingSample +from lczero_training.training.utils import make_weights_mask +from proto.training_config_pb2 import ( + LossConfig, + MovesLeftLossConfig, + PolicyLossConfig, + RegularizationLossConfig, + ValueCategoricalLossConfig, + ValueErrorLossConfig, + ValueLossConfig, +) + +from .model import LczeroModel, ModelPrediction + + +def _compute_q_from_wdl(wdl_logits: jax.Array) -> jax.Array: + """Compute Q value from WDL logits.""" + wdl_probs = jax.nn.softmax(wdl_logits) + q_weights = jnp.array([1.0, 0.0, -1.0]) + return jnp.dot(wdl_probs, q_weights) + + +class LossBase: + def __init__( + self, + config: Union[ + PolicyLossConfig, + ValueLossConfig, + MovesLeftLossConfig, + ValueErrorLossConfig, + ValueCategoricalLossConfig, + ], + ) -> None: + self.head_name = config.head_name + self.metric_name = config.metric_name or config.head_name + self.weight = config.weight + + def __call__( + self, + predictions: ModelPrediction, + sample: TrainingSample, + ) -> jax.Array: + raise NotImplementedError("Subclasses must implement __call__") + + +class RegularizationLoss: + """Computes regularization loss on model parameters.""" + + def __init__(self, config: RegularizationLossConfig) -> None: + self.metric_name = config.metric_name or "l2" + self.weight = config.weight + self._selector = config.selector + + def __call__(self, model: LczeroModel) -> jax.Array: + params = nnx.state(model, nnx.Param) + mask = make_weights_mask(self._selector, params) + masked_params = jax.tree.map( + lambda p, m: p.value if m else jnp.zeros_like(p.value), + params, + mask, + is_leaf=lambda x: isinstance(x, nnx.Variable), + ) + leaves = jax.tree.leaves(masked_params) + return sum( + (jnp.sum(jnp.square(leaf)) for leaf in leaves), jnp.array(0.0) + ) + + +class LczeroLoss: + policy_losses: List["PolicyLoss"] + value_losses: List["ValueLoss"] + movesleft_losses: List["MovesLeftLoss"] + value_error_losses: List["ValueErrorLoss"] + value_categorical_losses: List["ValueCategoricalLoss"] + regularization_losses: List["RegularizationLoss"] + + def __init__(self, config: LossConfig) -> None: + self.config = config + self.policy_losses = [ + PolicyLoss(loss_config) for loss_config in config.policy + ] + self.value_losses = [ + ValueLoss(loss_config) for loss_config in config.value + ] + self.movesleft_losses = [ + MovesLeftLoss(loss_config) for loss_config in config.movesleft + ] + self.value_error_losses = [ + ValueErrorLoss(loss_config) for loss_config in config.value_error + ] + self.value_categorical_losses = [ + ValueCategoricalLoss(loss_config) + for loss_config in config.value_categorical + ] + self.regularization_losses = [ + RegularizationLoss(loss_config) + for loss_config in config.regularization + ] + + def _validate_no_duplicate_metrics( + loss_type_name: str, + losses: Sequence[Union[LossBase, RegularizationLoss]], + ) -> None: + seen = set() + for name in (loss.metric_name for loss in losses): + if name in seen: + raise ValueError( + f"Duplicate metric name: {loss_type_name}/{name}" + ) + seen.add(name) + + _validate_no_duplicate_metrics("policy", self.policy_losses) + _validate_no_duplicate_metrics("value", self.value_losses) + _validate_no_duplicate_metrics("movesleft", self.movesleft_losses) + _validate_no_duplicate_metrics("value_error", self.value_error_losses) + _validate_no_duplicate_metrics( + "value_categorical", self.value_categorical_losses + ) + _validate_no_duplicate_metrics( + "regularization", self.regularization_losses + ) + + def __call__( + self, + model: LczeroModel, + sample: TrainingSample, + ) -> Tuple[jax.Array, Dict[str, jax.Array]]: + # Run model forward pass. + predictions = model(sample.inputs) + + unweighted_losses: Dict[str, jax.Array] = {} + weighted_losses: List[jax.Array] = [] + + for policy_loss in self.policy_losses: + loss = policy_loss(predictions, sample) + unweighted_losses[f"policy/{policy_loss.metric_name}"] = loss + weighted_losses.append(loss * policy_loss.weight) + + for value_loss in self.value_losses: + loss = value_loss(predictions, sample) + unweighted_losses[f"value/{value_loss.metric_name}"] = loss + weighted_losses.append(loss * value_loss.weight) + + for movesleft_loss in self.movesleft_losses: + loss = movesleft_loss(predictions, sample) + unweighted_losses[f"movesleft/{movesleft_loss.metric_name}"] = loss + weighted_losses.append(loss * movesleft_loss.weight) + + for value_error_loss in self.value_error_losses: + loss = value_error_loss(predictions, sample) + unweighted_losses[f"value_error/{value_error_loss.metric_name}"] = ( + loss + ) + weighted_losses.append(loss * value_error_loss.weight) + + for value_categorical_loss in self.value_categorical_losses: + loss = value_categorical_loss(predictions, sample) + unweighted_losses[ + f"value_categorical/{value_categorical_loss.metric_name}" + ] = loss + weighted_losses.append(loss * value_categorical_loss.weight) + + for reg_loss in self.regularization_losses: + loss = reg_loss(model) + unweighted_losses[f"regularization/{reg_loss.metric_name}"] = loss + weighted_losses.append(loss * reg_loss.weight) + + data_loss = jnp.sum(jnp.array(weighted_losses)) + + return data_loss, unweighted_losses + + +class ValueLoss(LossBase): + def __init__(self, config: ValueLossConfig) -> None: + super().__init__(config) + self.value_type = config.value_type + + def __call__( + self, + predictions: ModelPrediction, + sample: TrainingSample, + ) -> jax.Array: + value_pred = predictions.value[self.head_name] + value_logits = value_pred[0] + # Extract raw q/d from sample and compute WDL. + value_q = sample.values[self.value_type, 0] + value_d = sample.values[self.value_type, 1] + # Compute WDL: w = (1 + q - d) / 2, l = (1 - q - d) / 2 + value_w = (1.0 + value_q - value_d) / 2.0 + value_l = (1.0 - value_q - value_d) / 2.0 + value_wdl = jnp.stack([value_w, value_d, value_l], axis=-1) + + # The cross-entropy between the predicted value and the target value. + value_cross_entropy = optax.softmax_cross_entropy( + logits=value_logits, labels=jax.lax.stop_gradient(value_wdl) + ) + assert isinstance(value_cross_entropy, jax.Array) + return value_cross_entropy + + +class PolicyLoss(LossBase): + def __init__(self, config: PolicyLossConfig): + super().__init__(config) + self.config = config + if config.type == PolicyLossConfig.LOSS_TYPE_UNSPECIFIED: + raise ValueError( + f"Policy loss type must be specified for head '{config.head_name}'." + ) + self._loss_type = config.type + temperature = config.temperature + if temperature <= 0: + temperature = 1.0 + self._temperature = temperature + + # Store optimistic config if present. + if config.HasField("optimistic"): + opt = config.optimistic + self.opt_value_head: Optional[str] = opt.value_head_name + self.opt_value_type = opt.value_type + self.opt_strength = opt.strength + self.opt_eps = opt.eps + self.opt_alpha = opt.alpha + self.opt_propagate_gradients = opt.propagate_value_gradients + else: + self.opt_value_head = None + + def _apply_temperature_and_normalize( + self, policy_targets: jax.Array + ) -> jax.Array: + if self._temperature == 1.0: + return policy_targets + + # Apply temperature scaling. + policy_targets = jnp.power(policy_targets, 1.0 / self._temperature) + + # Renormalize after temperature scaling. + target_sum = jnp.sum(policy_targets, axis=-1, keepdims=True) + safe_sum = jnp.where( + target_sum > 0, target_sum, jnp.ones_like(target_sum) + ) + return policy_targets / safe_sum + + def _compute_optimistic_weight( + self, + value_pred: Tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]], + target_q: jax.Array, + ) -> jax.Array: + """Compute optimistic policy weight from value head predictions.""" + wdl_logits = value_pred[0] + error_pred = value_pred[1] + assert error_pred is not None, ( + "Error prediction required for optimistic weighting" + ) + + # Optionally block gradients to value and error heads. + if not self.opt_propagate_gradients: + wdl_logits = jax.lax.stop_gradient(wdl_logits) + error_pred = jax.lax.stop_gradient(error_pred) + + # Compute predicted Q from WDL. + q_pred = _compute_q_from_wdl(wdl_logits) + + # Compute sigma and z-score. + sigma = jnp.sqrt(error_pred.squeeze()) + z = (target_q - q_pred) / (sigma + self.opt_eps) + + # Compute weight. + return jax.nn.sigmoid((z - self.opt_strength) * self.opt_alpha) + + def __call__( + self, + predictions: ModelPrediction, + sample: TrainingSample, + ) -> jax.Array: + policy_pred = predictions.policy[self.head_name] + # Extract probabilities from sample. + policy_targets = jnp.asarray( + sample.probabilities, dtype=policy_pred.dtype + ) + if self.config.illegal_moves == PolicyLossConfig.MASK: + policy_pred = jnp.where(policy_targets >= 0, policy_pred, -jnp.inf) + + # Zero out negative targets for illegal moves. + policy_targets = jax.nn.relu(policy_targets) + + # Apply temperature scaling and renormalization if needed. + policy_targets = self._apply_temperature_and_normalize(policy_targets) + + cross_entropy = cast( + jax.Array, + optax.safe_softmax_cross_entropy( + logits=policy_pred, labels=policy_targets + ), + ) + if self._loss_type == PolicyLossConfig.CROSS_ENTROPY: + loss = cross_entropy + elif self._loss_type == PolicyLossConfig.KL: + loss = cross_entropy + jnp.sum( + xlogy(policy_targets, policy_targets), axis=-1 + ) + else: + raise AssertionError( + f"Unknown policy loss type: {self._loss_type}." + ) + + # Apply optimistic weighting if configured. + if self.opt_value_head is not None: + value_pred = predictions.value[self.opt_value_head] + target_q = sample.values[self.opt_value_type, 0] + loss = loss * self._compute_optimistic_weight(value_pred, target_q) + + return loss + + +class MovesLeftLoss(LossBase): + def __init__(self, config: MovesLeftLossConfig) -> None: + super().__init__(config) + self.value_type = config.value_type + + def __call__( + self, + predictions: ModelPrediction, + sample: TrainingSample, + ) -> jax.Array: + movesleft_pred = predictions.movesleft[self.head_name] + # Extract movesleft from sample. + # sample.values shape: [6, 3], component 2 is movesleft. + movesleft_targets = sample.values[self.value_type, 2] + + # Scale the loss to similar range as other losses. + scale = 20.0 + targets = movesleft_targets / scale + scaled_predictions = movesleft_pred / scale + + # Huber loss + huber_loss = optax.huber_loss( + predictions=scaled_predictions, targets=targets, delta=10.0 / scale + ) + assert isinstance(huber_loss, jax.Array) + return huber_loss.squeeze() + + +class ValueErrorLoss(LossBase): + def __init__(self, config: ValueErrorLossConfig) -> None: + super().__init__(config) + self.value_type = config.value_type + self.propagate_value_gradients = config.propagate_value_gradients + + def __call__( + self, + predictions: ModelPrediction, + sample: TrainingSample, + ) -> jax.Array: + value_pred = predictions.value[self.head_name] + wdl_logits = value_pred[0] + error_pred = value_pred[1] + assert error_pred is not None + + # Convert WDL to Q value. + predicted_q = _compute_q_from_wdl(wdl_logits) + + # Get target Q value. + target_q = sample.values[self.value_type, 0] + + # Compute actual squared error. + actual_squared_error = jnp.square(predicted_q - target_q) + + # Optionally block gradients to WDL head. + if not self.propagate_value_gradients: + actual_squared_error = jax.lax.stop_gradient(actual_squared_error) + + # MSE between error prediction and actual error. + loss = jnp.square(error_pred - actual_squared_error) + + return loss.squeeze() + + +class ValueCategoricalLoss(LossBase): + def __init__(self, config: ValueCategoricalLossConfig) -> None: + super().__init__(config) + self.value_type = config.value_type + + def __call__( + self, + predictions: ModelPrediction, + sample: TrainingSample, + ) -> jax.Array: + value_pred = predictions.value[self.head_name] + categorical_logits = value_pred[2] + assert categorical_logits is not None + + # Get target Q value from sample. + target_q = sample.values[self.value_type, 0] + + # Convert Q to bucket index: map [-1, 1) to [0, num_buckets). + num_buckets = categorical_logits.shape[-1] + bucket_index = jnp.floor((target_q + 1.0) / 2.0 * num_buckets).astype( + jnp.int32 + ) + bucket_index = jnp.clip(bucket_index, 0, num_buckets - 1) + + # Create one-hot target. + target_one_hot = jax.nn.one_hot(bucket_index, num_buckets) + + # Compute softmax cross-entropy. + loss = optax.softmax_cross_entropy( + logits=categorical_logits, + labels=jax.lax.stop_gradient(target_one_hot), + ) + assert isinstance(loss, jax.Array) + return loss diff --git a/src/lczero_training/model/model.py b/src/lczero_training/model/model.py new file mode 100644 index 00000000..1ce29eaf --- /dev/null +++ b/src/lczero_training/model/model.py @@ -0,0 +1,122 @@ +import dataclasses +import math +from typing import Optional, Tuple + +import jax +import jax.numpy as jnp +from flax import nnx + +from proto import model_config_pb2 + +from .embedding import Embedding +from .encoder import EncoderTower +from .movesleft_head import MovesLeftHead +from .policy_head import PolicyHead +from .utils import get_dtype +from .value_head import ValueHead + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass +class ModelPrediction: + """Output predictions from LczeroModel. + + Fields: + value: Dictionary mapping head names to value prediction tuples. + policy: Dictionary mapping head names to policy logits. + movesleft: Dictionary mapping head names to moves-left predictions. + """ + + value: dict[str, Tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]] + policy: dict[str, jax.Array] + movesleft: dict[str, jax.Array] + + +class LczeroModel(nnx.Module): + def __init__(self, config: model_config_pb2.ModelConfig, *, rngs: nnx.Rngs): + self.config = config + self._input_channels = 112 + deepnorm_beta = math.pow(8.0 * config.encoder.num_blocks, -0.25) + + self.embedding = Embedding( + input_channels=self._input_channels, + config=config.embedding, + defaults=config.defaults, + deepnorm_alpha=math.pow(2.0 * config.encoder.num_blocks, -0.25), + deepnorm_beta=deepnorm_beta, + rngs=rngs, + ) + + assert self.config.encoder.num_blocks > 0 + + self.encoders = EncoderTower( + in_features=config.embedding.embedding_size, + config=config.encoder, + defaults=config.defaults, + deepnorm_beta=deepnorm_beta, + rngs=rngs, + ) + + self.value_heads = nnx.Dict( + { + head_config.name: ValueHead( + in_features=config.embedding.embedding_size, + config=head_config, + defaults=config.defaults, + rngs=rngs, + ) + for head_config in config.value_head + } + ) + + # Named to appear before 'policy_heads' alphabetically in pytree state. + # This ensures shared embedding appears at parent level during + # serialization. + self.policy_embedding_shared: Optional[nnx.Linear] + if config.HasField("shared_policy_embedding_size"): + self.policy_embedding_shared = nnx.Linear( + in_features=config.embedding.embedding_size, + out_features=config.shared_policy_embedding_size, + rngs=rngs, + ) + else: + self.policy_embedding_shared = None + + self.policy_heads = nnx.Dict( + { + head_config.name: PolicyHead( + in_features=config.embedding.embedding_size, + config=head_config, + defaults=config.defaults, + shared_embedding=self.policy_embedding_shared, + rngs=rngs, + ) + for head_config in config.policy_head + } + ) + self.movesleft_heads = nnx.Dict( + { + head_config.name: MovesLeftHead( + in_features=config.embedding.embedding_size, + config=head_config, + defaults=config.defaults, + rngs=rngs, + ) + for head_config in config.movesleft_head + } + ) + + def __call__(self, x: jax.Array) -> ModelPrediction: + x = jnp.astype(x, get_dtype(self.config.defaults.compute_dtype)) + x = jnp.transpose(x, (1, 2, 0)) + x = jnp.reshape(x, (64, self._input_channels)) + x = self.embedding(x) + x = self.encoders(x) + + value = {name: head(x) for name, head in self.value_heads.items()} + policy = {name: head(x) for name, head in self.policy_heads.items()} + movesleft = { + name: head(x) for name, head in self.movesleft_heads.items() + } + + return ModelPrediction(value=value, policy=policy, movesleft=movesleft) diff --git a/src/lczero_training/model/movesleft_head.py b/src/lczero_training/model/movesleft_head.py new file mode 100644 index 00000000..9737bb2b --- /dev/null +++ b/src/lczero_training/model/movesleft_head.py @@ -0,0 +1,42 @@ +import jax +from flax import nnx + +from proto import model_config_pb2 + +from .utils import get_activation + + +class MovesLeftHead(nnx.Module): + def __init__( + self, + in_features: int, + config: model_config_pb2.MovesLeftHeadConfig, + defaults: model_config_pb2.DefaultsConfig, + *, + rngs: nnx.Rngs, + ): + self.activation = defaults.activation + self.embed = nnx.Linear( + in_features=in_features, + out_features=config.num_channels, + rngs=rngs, + ) + + self.dense1 = nnx.Linear( + in_features=config.num_channels * 64, + out_features=128, + rngs=rngs, + ) + self.out = nnx.Linear( + in_features=128, + out_features=1, + rngs=rngs, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.embed(x).flatten() + x = get_activation(self.activation)(x) + x = self.dense1(x) + x = get_activation(self.activation)(x) + x = self.out(x) + return nnx.relu(x) diff --git a/src/lczero_training/model/policy_head.py b/src/lczero_training/model/policy_head.py new file mode 100644 index 00000000..b2ac4c2f --- /dev/null +++ b/src/lczero_training/model/policy_head.py @@ -0,0 +1,234 @@ +import math +from typing import Optional + +import jax +import jax.numpy as jnp +from flax import nnx + +from proto import model_config_pb2 + +from .utils import get_activation # , get_policy_map + + +class PolicyHead(nnx.Module): + def __init__( + self, + in_features: int, + config: model_config_pb2.PolicyHeadConfig, + defaults: model_config_pb2.DefaultsConfig, + shared_embedding: Optional[nnx.Linear] = None, + *, + rngs: nnx.Rngs, + ): + assert (shared_embedding is not None) != config.HasField( + "embedding_size" + ) + self.activation = defaults.activation + if shared_embedding is not None: + self.tokens = shared_embedding + embedding_size = shared_embedding.out_features + else: + self.tokens = nnx.Linear( + in_features=in_features, + out_features=config.embedding_size, + rngs=rngs, + ) + embedding_size = config.embedding_size + + self.q = nnx.Linear( + in_features=embedding_size, + out_features=config.d_model, + rngs=rngs, + ) + + self.k = nnx.Linear( + in_features=embedding_size, + out_features=config.d_model, + rngs=rngs, + ) + + self.dk = math.sqrt(config.d_model) + self.promotion_dense = nnx.Linear( + in_features=config.d_model, + out_features=4, + use_bias=False, + rngs=rngs, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.tokens(x) + x = get_activation(self.activation)(x) + + q = self.q(x) + k = self.k(x) + qk = jnp.einsum("qd,kd->qk", q, k) + + promotion_keys = k[-8:, :] + promotion_offsets = self.promotion_dense(promotion_keys) + promotion_offsets = promotion_offsets.transpose((1, 0)) * self.dk + # knight offset is added to the other three + promotion_offsets = promotion_offsets[:3, :] + promotion_offsets[3:4, :] + + n_promo_logits = qk[-16:-8, -8:] + q_promo_logits = jnp.expand_dims( + n_promo_logits + promotion_offsets[0:1, :], axis=-1 + ) + r_promo_logits = jnp.expand_dims( + n_promo_logits + promotion_offsets[1:2, :], axis=-1 + ) + b_promo_logits = jnp.expand_dims( + n_promo_logits + promotion_offsets[2:3, :], axis=-1 + ) + promotion_logits = jnp.concatenate( + [q_promo_logits, r_promo_logits, b_promo_logits], axis=-1 + ) + policy_attn_logits = qk / self.dk + promotion_logits = promotion_logits.reshape((8, 24)) / self.dk + + logits = jnp.concatenate( + [policy_attn_logits.flatten(), promotion_logits.flatten()], axis=-1 + ) + + return logits[_policy_map] + + +# fmt: off +_policy_map = jnp.array([ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 16, 17, 18, 24, 27, 32, 36, 40, 45, 48, 54, 56, + 63, 64, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 80, 81, 82, 83, 89, 92, 97, + 101, 105, 110, 113, 119, 121, 128, 129, 131, 132, 133, 134, 135, 136, 137, 138, + 139, 140, 144, 145, 146, 147, 148, 154, 157, 162, 166, 170, 175, 178, 186, 192, + 193, 194, 196, 197, 198, 199, 201, 202, 203, 204, 205, 209, 210, 211, 212, 213, + 216, 219, 222, 227, 231, 235, 243, 251, 256, 257, 258, 259, 261, 262, 263, 266, + 267, 268, 269, 270, 274, 275, 276, 277, 278, 281, 284, 287, 288, 292, 300, 308, + 316, 320, 321, 322, 323, 324, 326, 327, 331, 332, 333, 334, 335, 339, 340, 341, + 342, 343, 346, 349, 353, 357, 360, 365, 373, 381, 384, 385, 386, 387, 388, 389, + 391, 396, 397, 398, 399, 404, 405, 406, 407, 411, 414, 418, 422, 425, 430, 432, + 438, 446, 448, 449, 450, 451, 452, 453, 454, 461, 462, 463, 469, 470, 471, 476, + 479, 483, 487, 490, 495, 497, 503, 504, 511, 512, 513, 514, 521, 522, 523, 524, + 525, 526, 527, 528, 529, 530, 536, 537, 538, 544, 547, 552, 556, 560, 565, 568, + 574, 576, 577, 578, 579, 584, 586, 587, 588, 589, 590, 591, 592, 593, 594, 595, + 600, 601, 602, 603, 609, 612, 617, 621, 625, 630, 633, 639, 640, 641, 642, 643, + 644, 648, 649, 651, 652, 653, 654, 655, 656, 657, 658, 659, 660, 664, 665, 666, + 667, 668, 674, 677, 682, 686, 690, 695, 698, 705, 706, 707, 708, 709, 712, 713, + 714, 716, 717, 718, 719, 721, 722, 723, 724, 725, 729, 730, 731, 732, 733, 736, + 739, 742, 747, 751, 755, 763, 770, 771, 772, 773, 774, 776, 777, 778, 779, 781, + 782, 783, 786, 787, 788, 789, 790, 794, 795, 796, 797, 798, 801, 804, 807, 808, + 812, 820, 828, 835, 836, 837, 838, 839, 840, 841, 842, 843, 844, 846, 847, 851, + 852, 853, 854, 855, 859, 860, 861, 862, 863, 866, 869, 873, 877, 880, 885, 893, + 900, 901, 902, 903, 904, 905, 906, 907, 908, 909, 911, 916, 917, 918, 919, 924, + 925, 926, 927, 931, 934, 938, 942, 945, 950, 952, 958, 965, 966, 967, 968, 969, + 970, 971, 972, 973, 974, 981, 982, 983, 989, 990, 991, 996, 999, 1003, 1007, + 1010, 1015, 1017, 1023, 1024, 1025, 1026, 1032, 1033, 1034, 1041, 1042, 1043, + 1044, 1045, 1046, 1047, 1048, 1049, 1050, 1056, 1057, 1058, 1064, 1067, 1072, + 1076, 1080, 1085, 1088, 1089, 1090, 1091, 1096, 1097, 1098, 1099, 1104, 1106, + 1107, 1108, 1109, 1110, 1111, 1112, 1113, 1114, 1115, 1120, 1121, 1122, 1123, + 1129, 1132, 1137, 1141, 1145, 1150, 1152, 1153, 1154, 1155, 1156, 1160, 1161, + 1162, 1163, 1164, 1168, 1169, 1171, 1172, 1173, 1174, 1175, 1176, 1177, 1178, + 1179, 1180, 1184, 1185, 1186, 1187, 1188, 1194, 1197, 1202, 1206, 1210, 1215, + 1217, 1218, 1219, 1220, 1221, 1225, 1226, 1227, 1228, 1229, 1232, 1233, 1234, + 1236, 1237, 1238, 1239, 1241, 1242, 1243, 1244, 1245, 1249, 1250, 1251, 1252, + 1253, 1256, 1259, 1262, 1267, 1271, 1275, 1282, 1283, 1284, 1285, 1286, 1290, + 1291, 1292, 1293, 1294, 1296, 1297, 1298, 1299, 1301, 1302, 1303, 1306, 1307, + 1308, 1309, 1310, 1314, 1315, 1316, 1317, 1318, 1321, 1324, 1327, 1328, 1332, + 1340, 1347, 1348, 1349, 1350, 1351, 1355, 1356, 1357, 1358, 1359, 1360, 1361, + 1362, 1363, 1364, 1366, 1367, 1371, 1372, 1373, 1374, 1375, 1379, 1380, 1381, + 1382, 1383, 1386, 1389, 1393, 1397, 1400, 1405, 1412, 1413, 1414, 1415, 1420, + 1421, 1422, 1423, 1424, 1425, 1426, 1427, 1428, 1429, 1431, 1436, 1437, 1438, + 1439, 1444, 1445, 1446, 1447, 1451, 1454, 1458, 1462, 1465, 1470, 1477, 1478, + 1479, 1485, 1486, 1487, 1488, 1489, 1490, 1491, 1492, 1493, 1494, 1501, 1502, + 1503, 1509, 1510, 1511, 1516, 1519, 1523, 1527, 1530, 1535, 1536, 1539, 1544, + 1545, 1546, 1552, 1553, 1554, 1561, 1562, 1563, 1564, 1565, 1566, 1567, 1568, + 1569, 1570, 1576, 1577, 1578, 1584, 1587, 1592, 1596, 1601, 1604, 1608, 1609, + 1610, 1611, 1616, 1617, 1618, 1619, 1624, 1626, 1627, 1628, 1629, 1630, 1631, + 1632, 1633, 1634, 1635, 1640, 1641, 1642, 1643, 1649, 1652, 1657, 1661, 1666, + 1669, 1672, 1673, 1674, 1675, 1676, 1680, 1681, 1682, 1683, 1684, 1688, 1689, + 1691, 1692, 1693, 1694, 1695, 1696, 1697, 1698, 1699, 1700, 1704, 1705, 1706, + 1707, 1708, 1714, 1717, 1722, 1726, 1728, 1731, 1734, 1737, 1738, 1739, 1740, + 1741, 1745, 1746, 1747, 1748, 1749, 1752, 1753, 1754, 1756, 1757, 1758, 1759, + 1761, 1762, 1763, 1764, 1765, 1769, 1770, 1771, 1772, 1773, 1776, 1779, 1782, + 1787, 1791, 1793, 1796, 1799, 1802, 1803, 1804, 1805, 1806, 1810, 1811, 1812, + 1813, 1814, 1816, 1817, 1818, 1819, 1821, 1822, 1823, 1826, 1827, 1828, 1829, + 1830, 1834, 1835, 1836, 1837, 1838, 1841, 1844, 1847, 1848, 1852, 1858, 1861, + 1867, 1868, 1869, 1870, 1871, 1875, 1876, 1877, 1878, 1879, 1880, 1881, 1882, + 1883, 1884, 1886, 1887, 1891, 1892, 1893, 1894, 1895, 1899, 1900, 1901, 1902, + 1903, 1906, 1909, 1913, 1917, 1923, 1926, 1932, 1933, 1934, 1935, 1940, 1941, + 1942, 1943, 1944, 1945, 1946, 1947, 1948, 1949, 1951, 1956, 1957, 1958, 1959, + 1964, 1965, 1966, 1967, 1971, 1974, 1978, 1982, 1988, 1991, 1997, 1998, 1999, + 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2021, 2022, 2023, + 2029, 2030, 2031, 2036, 2039, 2043, 2047, 2048, 2052, 2056, 2059, 2064, 2065, + 2066, 2072, 2073, 2074, 2081, 2082, 2083, 2084, 2085, 2086, 2087, 2088, 2089, + 2090, 2096, 2097, 2098, 2104, 2107, 2113, 2117, 2121, 2124, 2128, 2129, 2130, + 2131, 2136, 2137, 2138, 2139, 2144, 2146, 2147, 2148, 2149, 2150, 2151, 2152, + 2153, 2154, 2155, 2160, 2161, 2162, 2163, 2169, 2172, 2178, 2182, 2186, 2189, + 2192, 2193, 2194, 2195, 2196, 2200, 2201, 2202, 2203, 2204, 2208, 2209, 2211, + 2212, 2213, 2214, 2215, 2216, 2217, 2218, 2219, 2220, 2224, 2225, 2226, 2227, + 2228, 2234, 2237, 2243, 2247, 2248, 2251, 2254, 2257, 2258, 2259, 2260, 2261, + 2265, 2266, 2267, 2268, 2269, 2272, 2273, 2274, 2276, 2277, 2278, 2279, 2281, + 2282, 2283, 2284, 2285, 2289, 2290, 2291, 2292, 2293, 2296, 2299, 2302, 2304, + 2308, 2313, 2316, 2319, 2322, 2323, 2324, 2325, 2326, 2330, 2331, 2332, 2333, + 2334, 2336, 2337, 2338, 2339, 2341, 2342, 2343, 2346, 2347, 2348, 2349, 2350, + 2354, 2355, 2356, 2357, 2358, 2361, 2364, 2367, 2369, 2373, 2378, 2381, 2387, + 2388, 2389, 2390, 2391, 2395, 2396, 2397, 2398, 2399, 2400, 2401, 2402, 2403, + 2404, 2406, 2407, 2411, 2412, 2413, 2414, 2415, 2419, 2420, 2421, 2422, 2423, + 2426, 2429, 2434, 2438, 2443, 2446, 2452, 2453, 2454, 2455, 2460, 2461, 2462, + 2463, 2464, 2465, 2466, 2467, 2468, 2469, 2471, 2476, 2477, 2478, 2479, 2484, + 2485, 2486, 2487, 2491, 2494, 2499, 2503, 2508, 2511, 2517, 2518, 2519, 2525, + 2526, 2527, 2528, 2529, 2530, 2531, 2532, 2533, 2534, 2541, 2542, 2543, 2549, + 2550, 2551, 2556, 2559, 2560, 2565, 2568, 2572, 2576, 2579, 2584, 2585, 2586, + 2592, 2593, 2594, 2601, 2602, 2603, 2604, 2605, 2606, 2607, 2608, 2609, 2610, + 2616, 2617, 2618, 2625, 2630, 2633, 2637, 2641, 2644, 2648, 2649, 2650, 2651, + 2656, 2657, 2658, 2659, 2664, 2666, 2667, 2668, 2669, 2670, 2671, 2672, 2673, + 2674, 2675, 2680, 2681, 2682, 2683, 2690, 2695, 2698, 2702, 2706, 2709, 2712, + 2713, 2714, 2715, 2716, 2720, 2721, 2722, 2723, 2724, 2728, 2729, 2731, 2732, + 2733, 2734, 2735, 2736, 2737, 2738, 2739, 2740, 2744, 2745, 2746, 2747, 2748, + 2755, 2763, 2767, 2768, 2771, 2774, 2777, 2778, 2779, 2780, 2781, 2785, 2786, + 2787, 2788, 2789, 2792, 2793, 2794, 2796, 2797, 2798, 2799, 2801, 2802, 2803, + 2804, 2805, 2809, 2810, 2811, 2812, 2813, 2820, 2824, 2828, 2833, 2836, 2839, + 2842, 2843, 2844, 2845, 2846, 2850, 2851, 2852, 2853, 2854, 2856, 2857, 2858, + 2859, 2861, 2862, 2863, 2866, 2867, 2868, 2869, 2870, 2874, 2875, 2876, 2877, + 2878, 2880, 2885, 2889, 2893, 2898, 2901, 2907, 2908, 2909, 2910, 2911, 2915, + 2916, 2917, 2918, 2919, 2920, 2921, 2922, 2923, 2924, 2926, 2927, 2931, 2932, + 2933, 2934, 2935, 2939, 2940, 2941, 2942, 2943, 2945, 2950, 2954, 2958, 2963, + 2966, 2972, 2973, 2974, 2975, 2980, 2981, 2982, 2983, 2984, 2985, 2986, 2987, + 2988, 2989, 2991, 2996, 2997, 2998, 2999, 3004, 3005, 3006, 3007, 3010, 3015, + 3019, 3023, 3028, 3031, 3037, 3038, 3039, 3045, 3046, 3047, 3048, 3049, 3050, + 3051, 3052, 3053, 3054, 3061, 3062, 3063, 3069, 3070, 3071, 3072, 3078, 3080, + 3085, 3088, 3092, 3096, 3099, 3104, 3105, 3106, 3112, 3113, 3114, 3121, 3122, + 3123, 3124, 3125, 3126, 3127, 3128, 3129, 3130, 3137, 3143, 3145, 3150, 3153, + 3157, 3161, 3164, 3168, 3169, 3170, 3171, 3176, 3177, 3178, 3179, 3184, 3186, + 3187, 3188, 3189, 3190, 3191, 3192, 3193, 3194, 3195, 3202, 3210, 3215, 3218, + 3222, 3226, 3229, 3232, 3233, 3234, 3235, 3236, 3240, 3241, 3242, 3243, 3244, + 3248, 3249, 3251, 3252, 3253, 3254, 3255, 3256, 3257, 3258, 3259, 3260, 3267, + 3275, 3283, 3287, 3288, 3291, 3294, 3297, 3298, 3299, 3300, 3301, 3305, 3306, + 3307, 3308, 3309, 3312, 3313, 3314, 3316, 3317, 3318, 3319, 3321, 3322, 3323, + 3324, 3325, 3332, 3340, 3344, 3348, 3353, 3356, 3359, 3362, 3363, 3364, 3365, + 3366, 3370, 3371, 3372, 3373, 3374, 3376, 3377, 3378, 3379, 3381, 3382, 3383, + 3386, 3387, 3388, 3389, 3390, 3397, 3400, 3405, 3409, 3413, 3418, 3421, 3427, + 3428, 3429, 3430, 3431, 3435, 3436, 3437, 3438, 3439, 3440, 3441, 3442, 3443, + 3444, 3446, 3447, 3451, 3452, 3453, 3454, 3455, 3456, 3462, 3465, 3470, 3474, + 3478, 3483, 3486, 3492, 3493, 3494, 3495, 3500, 3501, 3502, 3503, 3504, 3505, + 3506, 3507, 3508, 3509, 3511, 3516, 3517, 3518, 3519, 3521, 3527, 3530, 3535, + 3539, 3543, 3548, 3551, 3557, 3558, 3559, 3565, 3566, 3567, 3568, 3569, 3570, + 3571, 3572, 3573, 3574, 3581, 3582, 3583, 3584, 3591, 3592, 3598, 3600, 3605, + 3608, 3612, 3616, 3619, 3624, 3625, 3626, 3632, 3633, 3634, 3641, 3642, 3643, + 3644, 3645, 3646, 3647, 3649, 3657, 3663, 3665, 3670, 3673, 3677, 3681, 3684, + 3688, 3689, 3690, 3691, 3696, 3697, 3698, 3699, 3704, 3706, 3707, 3708, 3709, + 3710, 3711, 3714, 3722, 3730, 3735, 3738, 3742, 3746, 3749, 3752, 3753, 3754, + 3755, 3756, 3760, 3761, 3762, 3763, 3764, 3768, 3769, 3771, 3772, 3773, 3774, + 3775, 3779, 3787, 3795, 3803, 3807, 3808, 3811, 3814, 3817, 3818, 3819, 3820, + 3821, 3825, 3826, 3827, 3828, 3829, 3832, 3833, 3834, 3836, 3837, 3838, 3839, + 3844, 3852, 3860, 3864, 3868, 3873, 3876, 3879, 3882, 3883, 3884, 3885, 3886, + 3890, 3891, 3892, 3893, 3894, 3896, 3897, 3898, 3899, 3901, 3902, 3903, 3909, + 3917, 3920, 3925, 3929, 3933, 3938, 3941, 3947, 3948, 3949, 3950, 3951, 3955, + 3956, 3957, 3958, 3959, 3960, 3961, 3962, 3963, 3964, 3966, 3967, 3974, 3976, + 3982, 3985, 3990, 3994, 3998, 4003, 4006, 4012, 4013, 4014, 4015, 4020, 4021, + 4022, 4023, 4024, 4025, 4026, 4027, 4028, 4029, 4031, 4032, 4039, 4041, 4047, + 4050, 4055, 4059, 4063, 4068, 4071, 4077, 4078, 4079, 4085, 4086, 4087, 4088, + 4089, 4090, 4091, 4092, 4093, 4094, 4096, 4097, 4098, 4099, 4100, 4101, 4120, + 4121, 4122, 4123, 4124, 4125, 4126, 4127, 4128, 4147, 4148, 4149, 4150, 4151, + 4152, 4153, 4154, 4155, 4174, 4175, 4176, 4177, 4178, 4179, 4180, 4181, 4182, + 4201, 4202, 4203, 4204, 4205, 4206, 4207, 4208, 4209, 4228, 4229, 4230, 4231, + 4232, 4233, 4234, 4235, 4236, 4255, 4256, 4257, 4258, 4259, 4260, 4261, 4262, + 4263, 4282, 4283, 4284, 4285, 4286, 4287, +], jnp.int32) diff --git a/src/lczero_training/model/shared.py b/src/lczero_training/model/shared.py new file mode 100644 index 00000000..e55e5392 --- /dev/null +++ b/src/lczero_training/model/shared.py @@ -0,0 +1,44 @@ +import jax +from flax import nnx +from flax.linen import initializers as flax_initializers + +from proto import net_pb2 + +from .utils import get_activation + + +class Ffn(nnx.Module): + def __init__( + self, + in_features: int, + hidden_features: int, + hidden_activation: net_pb2.NetworkFormat.ActivationFunction, + deepnorm_beta: float, + *, + rngs: nnx.Rngs, + ): + deepnorm_init = flax_initializers.variance_scaling( + scale=deepnorm_beta, + mode="fan_avg", + distribution="truncated_normal", + ) + out_features = in_features + self.linear1 = nnx.Linear( + in_features=in_features, + out_features=hidden_features, + kernel_init=deepnorm_init, + rngs=rngs, + ) + self.activation = hidden_activation + self.linear2 = nnx.Linear( + in_features=hidden_features, + out_features=out_features, + kernel_init=deepnorm_init, + rngs=rngs, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.linear1(x) + x = get_activation(self.activation)(x) + x = self.linear2(x) + return x diff --git a/src/lczero_training/model/utils.py b/src/lczero_training/model/utils.py new file mode 100644 index 00000000..ebe69f17 --- /dev/null +++ b/src/lczero_training/model/utils.py @@ -0,0 +1,50 @@ +from typing import Any + +import jax.numpy as jnp +from flax import nnx +from jax.nn import mish + +from proto import net_pb2 +from proto.hlo_pb2 import XlaShapeProto + + +def get_activation( + activation: net_pb2.NetworkFormat.ActivationFunction, +) -> Any: + return { + net_pb2.NetworkFormat.ACTIVATION_MISH: mish, + net_pb2.NetworkFormat.ACTIVATION_RELU: nnx.relu, + net_pb2.NetworkFormat.ACTIVATION_NONE: lambda x: x, + net_pb2.NetworkFormat.ACTIVATION_TANH: nnx.tanh, + net_pb2.NetworkFormat.ACTIVATION_SIGMOID: nnx.sigmoid, + net_pb2.NetworkFormat.ACTIVATION_SELU: nnx.selu, + net_pb2.NetworkFormat.ACTIVATION_SWISH: nnx.swish, + net_pb2.NetworkFormat.ACTIVATION_SOFTMAX: nnx.softmax, + }[activation] + + +def get_dtype(dtype: XlaShapeProto.Type) -> jnp.dtype: + return { + XlaShapeProto.PRED: jnp.bool_, + XlaShapeProto.S4: jnp.int4, + XlaShapeProto.S8: jnp.int8, + XlaShapeProto.S16: jnp.int16, + XlaShapeProto.S32: jnp.int32, + XlaShapeProto.S64: jnp.int64, + XlaShapeProto.U4: jnp.uint4, + XlaShapeProto.U8: jnp.uint8, + XlaShapeProto.U16: jnp.uint16, + XlaShapeProto.U32: jnp.uint32, + XlaShapeProto.U64: jnp.uint64, + XlaShapeProto.F16: jnp.float16, + XlaShapeProto.F32: jnp.float32, + XlaShapeProto.BF16: jnp.bfloat16, + XlaShapeProto.F64: jnp.float64, + XlaShapeProto.F8E5M2: jnp.float8_e5m2, + XlaShapeProto.F8E4M3FN: jnp.float8_e4m3fn, + XlaShapeProto.F8E4M3B11FNUZ: jnp.float8_e4m3b11fnuz, + XlaShapeProto.F8E5M2FNUZ: jnp.float8_e5m2fnuz, + XlaShapeProto.F8E4M3FNUZ: jnp.float8_e4m3fnuz, + XlaShapeProto.C64: jnp.complex64, + XlaShapeProto.C128: jnp.complex128, + }[dtype] diff --git a/src/lczero_training/model/value_head.py b/src/lczero_training/model/value_head.py new file mode 100644 index 00000000..48fcb040 --- /dev/null +++ b/src/lczero_training/model/value_head.py @@ -0,0 +1,67 @@ +from typing import Optional, Tuple + +import jax +from flax import nnx + +from proto import model_config_pb2 + +from .utils import get_activation + + +class ValueHead(nnx.Module): + def __init__( + self, + in_features: int, + config: model_config_pb2.ValueHeadConfig, + defaults: model_config_pb2.DefaultsConfig, + rngs: nnx.Rngs, + ): + self.activation = defaults.activation + self.has_error_output = config.has_error_output + self.num_categorical_buckets = config.num_categorical_buckets + self.embed = nnx.Linear( + in_features=in_features, + out_features=config.num_channels, + rngs=rngs, + ) + self.dense1 = nnx.Linear( + in_features=config.num_channels * 64, + out_features=128, + rngs=rngs, + ) + self.wdl = nnx.Linear( + in_features=128, + out_features=3, + rngs=rngs, + ) + if self.has_error_output: + self.error = nnx.Linear( + in_features=128, + out_features=1, + rngs=rngs, + ) + if self.num_categorical_buckets > 0: + self.categorical = nnx.Linear( + in_features=128, + out_features=self.num_categorical_buckets, + rngs=rngs, + ) + + def __call__( + self, x: jax.Array + ) -> Tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]: + x = self.embed(x).flatten() + x = get_activation(self.activation)(x) + x = self.dense1(x) + x = get_activation(self.activation)(x) + + wdl = self.wdl(x) + error = nnx.sigmoid(self.error(x)) if self.has_error_output else None + categorical = ( + self.categorical(x) if self.num_categorical_buckets > 0 else None + ) + + return (wdl, error, categorical) + + def predict(self, x: jax.Array) -> jax.Array: + return nnx.softmax(self(x)[0]) diff --git a/src/lczero_training/py.typed b/src/lczero_training/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/src/lczero_training/tests/test_protobuf.py b/src/lczero_training/tests/test_protobuf.py new file mode 100644 index 00000000..7fe2739e --- /dev/null +++ b/src/lczero_training/tests/test_protobuf.py @@ -0,0 +1,67 @@ +"""Test protobuf compilation and functionality.""" + + +def test_protobuf_import() -> None: + """Test that protobuf files can be imported.""" + import proto.data_loader_config_pb2 as data_loader_config_pb2 + import proto.model_config_pb2 as model_config_pb2 + import proto.root_config_pb2 as root_config_pb2 + import proto.training_config_pb2 as training_config_pb2 + import proto.training_metrics_pb2 as training_metrics_pb2 + + # Test creating config objects + data_loader_config = data_loader_config_pb2.DataLoaderConfig() + assert data_loader_config is not None + + root_config = root_config_pb2.RootConfig() + assert root_config is not None + + training_config = training_config_pb2.TrainingConfig() + assert training_config is not None + + model_config = model_config_pb2.ModelConfig() + assert model_config is not None + + metrics = training_metrics_pb2.DataLoaderMetricsProto() + assert metrics is not None + + +def test_protobuf_functionality() -> None: + """Test basic protobuf functionality.""" + import proto.data_loader_config_pb2 as data_loader_config_pb2 + import proto.root_config_pb2 as root_config_pb2 + + # Create a config and set some values + config = data_loader_config_pb2.DataLoaderConfig() + stage = config.stage.add() + stage.name = "file_path_provider" + stage.file_path_provider.directory = "/test/path" + + # Serialize and deserialize + serialized = config.SerializeToString() + assert len(serialized) > 0 + + config2 = data_loader_config_pb2.DataLoaderConfig() + config2.ParseFromString(serialized) + + assert config2.stage[0].file_path_provider.directory == "/test/path" + + # Test RootConfig functionality + root_config = root_config_pb2.RootConfig() + root_config.name = "test_config" + stage_config = root_config.data_loader.stage.add() + stage_config.name = "file_path_provider" + stage_config.file_path_provider.directory = "/test/path" + + # Serialize and deserialize root config + root_serialized = root_config.SerializeToString() + assert len(root_serialized) > 0 + + root_config2 = root_config_pb2.RootConfig() + root_config2.ParseFromString(root_serialized) + + assert root_config2.name == "test_config" + assert ( + root_config2.data_loader.stage[0].file_path_provider.directory + == "/test/path" + ) diff --git a/src/lczero_training/tests/test_protocol_registry.py b/src/lczero_training/tests/test_protocol_registry.py new file mode 100644 index 00000000..8d4c6cbd --- /dev/null +++ b/src/lczero_training/tests/test_protocol_registry.py @@ -0,0 +1,133 @@ +"""Test script for the protocol registry system.""" + +from dataclasses import dataclass +from typing import Any + +import pytest + +from lczero_training.daemon.protocol.registry import ( + CLASS_TO_TYPE_MAP, + TYPE_TO_CLASS_MAP, + register, +) + + +@pytest.fixture(autouse=True) +def clear_registry() -> Any: + """Clear registry maps before each test.""" + TYPE_TO_CLASS_MAP.clear() + CLASS_TO_TYPE_MAP.clear() + yield + TYPE_TO_CLASS_MAP.clear() + CLASS_TO_TYPE_MAP.clear() + + +def test_basic_registration() -> None: + """Test basic event type registration.""" + + @register("test_event") + @dataclass + class BasicPayload: + content: str + + # Check forward mapping + assert TYPE_TO_CLASS_MAP["test_event"] == BasicPayload + # Check reverse mapping + assert CLASS_TO_TYPE_MAP[BasicPayload] == "test_event" + + +def test_duplicate_event_type() -> None: + """Test that duplicate event types are rejected.""" + + @register("duplicate_event") + @dataclass + class FirstPayload: + data: str + + with pytest.raises( + ValueError, match=r".*duplicate_event.*already registered.*" + ): + + @register("duplicate_event") # Should fail + @dataclass + class SecondPayload: + other_data: int + + +def test_duplicate_class() -> None: + """Test that duplicate classes are rejected.""" + + @dataclass + class PayloadClass: + data: str + + # Register once + register("first_event")(PayloadClass) + + # Try to register same class again - should fail + with pytest.raises( + ValueError, match=r".*PayloadClass.*already registered.*" + ): + register("second_event")(PayloadClass) + + +def test_non_class_registration() -> None: + """Test that non-classes are rejected.""" + with pytest.raises(TypeError, match=r".*can only be used on classes.*"): + # Try to register a string instead of a class + @register("invalid_event") # type: ignore[arg-type] + def not_a_class() -> None: + pass + + +def test_multiple_registrations() -> None: + """Test multiple valid registrations work correctly.""" + + @register("event_one") + @dataclass + class PayloadOne: + data: str + + @register("event_two") + @dataclass + class PayloadTwo: + value: int + + @register("event_three") + @dataclass + class PayloadThree: + items: list + + # Check all mappings exist + assert TYPE_TO_CLASS_MAP["event_one"] == PayloadOne + assert TYPE_TO_CLASS_MAP["event_two"] == PayloadTwo + assert TYPE_TO_CLASS_MAP["event_three"] == PayloadThree + + assert CLASS_TO_TYPE_MAP[PayloadOne] == "event_one" + assert CLASS_TO_TYPE_MAP[PayloadTwo] == "event_two" + assert CLASS_TO_TYPE_MAP[PayloadThree] == "event_three" + + # Check we have exactly 3 entries in each map + assert len(TYPE_TO_CLASS_MAP) == 3 + assert len(CLASS_TO_TYPE_MAP) == 3 + + +def test_registry_persistence() -> None: + """Test that registry persists across imports.""" + + @register("persistent_event") + @dataclass + class PersistentPayload: + data: str + + # Re-import the module + from lczero_training.daemon.protocol.registry import ( + CLASS_TO_TYPE_MAP as imported_class_map, + ) + from lczero_training.daemon.protocol.registry import ( + TYPE_TO_CLASS_MAP as imported_type_map, + ) + + # Check the registration persists + assert imported_type_map["persistent_event"] == PersistentPayload + assert imported_class_map[PersistentPayload] == "persistent_event" diff --git a/src/lczero_training/tests/test_weights_tool.py b/src/lczero_training/tests/test_weights_tool.py new file mode 100644 index 00000000..e9e22654 --- /dev/null +++ b/src/lczero_training/tests/test_weights_tool.py @@ -0,0 +1,237 @@ +"""Test weights tool arithmetic operations.""" + +import os +import tempfile + +import numpy as np + +import proto.net_pb2 as net_pb2 +from lczero_training.tools.weight_wrappers import NetWrapper +from lczero_training.tools.weights_tool import load_weights, save_weights + + +def test_weights_arithmetic() -> None: + """Test arithmetic operations on simple networks.""" + # Create network A with a single layer containing value 10.0. + net_a = net_pb2.Net() + net_a.format.weights_encoding = net_pb2.Format.LINEAR16 + net_a.weights.ip1_val_w.min_val = 10.0 + net_a.weights.ip1_val_w.max_val = 10.0 + # LINEAR16 encoding: value 10.0 maps to uint16 value 32767 (mid-point). + net_a.weights.ip1_val_w.params = np.array( + [32767], dtype=np.uint16 + ).tobytes() + + # Create network B with a single layer containing value 20.0. + net_b = net_pb2.Net() + net_b.format.weights_encoding = net_pb2.Format.LINEAR16 + net_b.weights.ip1_val_w.min_val = 20.0 + net_b.weights.ip1_val_w.max_val = 20.0 + net_b.weights.ip1_val_w.params = np.array( + [32767], dtype=np.uint16 + ).tobytes() + + # Wrap the networks. + wrapper_a = NetWrapper(net_a) + wrapper_b = NetWrapper(net_b) + + # Compute: output = 0.2*A + 0.8*B + # Expected: 0.2*10 + 0.8*20 = 2 + 16 = 18 + output = 0.2 * wrapper_a + 0.8 * wrapper_b + + # Check the result. + result_value = output.weights.ip1_val_w.value + expected = 18.0 + + assert result_value.shape == (1,), ( + f"Expected shape (1,), got {result_value.shape}" + ) + assert np.isclose(result_value[0], expected, rtol=1e-4), ( + f"Expected {expected}, got {result_value[0]}" + ) + + +def test_policy_head_replacement() -> None: + """Test that assigning policy_heads actually replaces the data.""" + # Create network A with policy head value 10.0. + net_a = net_pb2.Net() + net_a.format.weights_encoding = net_pb2.Format.LINEAR16 + net_a.weights.policy_heads.ip_pol_w.min_val = 10.0 + net_a.weights.policy_heads.ip_pol_w.max_val = 10.0 + net_a.weights.policy_heads.ip_pol_w.params = np.array( + [32767], dtype=np.uint16 + ).tobytes() + + # Create network B with policy head value 20.0. + net_b = net_pb2.Net() + net_b.format.weights_encoding = net_pb2.Format.LINEAR16 + net_b.weights.policy_heads.ip_pol_w.min_val = 20.0 + net_b.weights.policy_heads.ip_pol_w.max_val = 20.0 + net_b.weights.policy_heads.ip_pol_w.params = np.array( + [32767], dtype=np.uint16 + ).tobytes() + + # Wrap and perform assignment. + wrapper_a = NetWrapper(net_a) + wrapper_b = NetWrapper(net_b) + + # Replace A's policy heads with B's. + wrapper_a.weights.policy_heads = wrapper_b.weights.policy_heads + + # Verify in-memory replacement. + result_value = wrapper_a.weights.policy_heads.ip_pol_w.value + assert np.isclose(result_value[0], 20.0, rtol=1e-4), ( + f"Expected 20.0 (B's value), got {result_value[0]}" + ) + + # Verify persistence (round-trip). + with tempfile.NamedTemporaryFile(suffix=".pb.gz", delete=False) as tmp: + tmp_path = tmp.name + try: + save_weights(wrapper_a, tmp_path) + reloaded = load_weights(tmp_path) + reloaded_value = reloaded.weights.policy_heads.ip_pol_w.value + assert np.isclose(reloaded_value[0], 20.0, rtol=1e-4), ( + f"After save/load: Expected 20.0, got {reloaded_value[0]}" + ) + finally: + os.unlink(tmp_path) + + +def test_policy_head_map_assignment() -> None: + """Test assigning policy_heads when one has policy_head_map and another doesn't.""" + # Create network A with NO policy_head_map. + net_a = net_pb2.Net() + net_a.format.weights_encoding = net_pb2.Format.LINEAR16 + net_a.weights.policy_heads.ip_pol_w.min_val = 5.0 + net_a.weights.policy_heads.ip_pol_w.max_val = 5.0 + net_a.weights.policy_heads.ip_pol_w.params = np.array( + [32767], dtype=np.uint16 + ).tobytes() + + # Create network B WITH policy_head_map. + net_b = net_pb2.Net() + net_b.format.weights_encoding = net_pb2.Format.LINEAR16 + # Add a policy_head_map entry with required key and value fields. + policy_map = net_b.weights.policy_heads.policy_head_map.add() + policy_map.key = "test_policy" + policy_map.value.ip_pol_w.min_val = 15.0 + policy_map.value.ip_pol_w.max_val = 15.0 + policy_map.value.ip_pol_w.params = np.array( + [32767], dtype=np.uint16 + ).tobytes() + + # Wrap networks. + wrapper_a = NetWrapper(net_a) + wrapper_b = NetWrapper(net_b) + + # Replace A's policy heads with B's (which has policy_head_map). + wrapper_a.weights.policy_heads = wrapper_b.weights.policy_heads + + # Access policy_head_map through the cache to verify consistency. + assert len(wrapper_a.weights.policy_heads.policy_head_map) == 1 + assert ( + wrapper_a.weights.policy_heads.policy_head_map[0]._proto.key + == "test_policy" + ) + + # Verify can save without serialization errors. + with tempfile.NamedTemporaryFile(suffix=".pb.gz", delete=False) as tmp: + tmp_path = tmp.name + try: + save_weights(wrapper_a, tmp_path) + # Verify round-trip preserves the policy_head_map. + reloaded = load_weights(tmp_path) + assert len(reloaded.weights.policy_heads.policy_head_map) == 1 + assert ( + reloaded.weights.policy_heads.policy_head_map[0]._proto.key + == "test_policy" + ) + policy_value = reloaded.weights.policy_heads.policy_head_map[ + 0 + ].value.ip_pol_w.value + assert np.isclose(policy_value[0], 15.0, rtol=1e-4) + finally: + os.unlink(tmp_path) + + +def test_noop_arithmetic() -> None: + """Test that 0.3*A + 0.7*A equals A (no-op operation).""" + # Create network A with simple weights (value 10.0). + net_a = net_pb2.Net() + net_a.format.weights_encoding = net_pb2.Format.LINEAR16 + net_a.weights.ip1_val_w.min_val = 10.0 + net_a.weights.ip1_val_w.max_val = 10.0 + net_a.weights.ip1_val_w.params = np.array( + [32767], dtype=np.uint16 + ).tobytes() + + wrapper_a = NetWrapper(net_a) + + # Perform no-op operation: 0.3*A + 0.7*A should equal A. + result = 0.3 * wrapper_a + 0.7 * wrapper_a + + # Verify in-memory: result should equal A (10.0). + result_value = result.weights.ip1_val_w.value + expected = 10.0 + assert np.isclose(result_value[0], expected, rtol=1e-4), ( + f"Expected {expected}, got {result_value[0]}" + ) + + # Verify persistence: save and reload. + with tempfile.NamedTemporaryFile(suffix=".pb.gz", delete=False) as tmp: + tmp_path = tmp.name + try: + save_weights(result, tmp_path) + reloaded = load_weights(tmp_path) + reloaded_value = reloaded.weights.ip1_val_w.value + assert np.isclose(reloaded_value[0], expected, rtol=1e-4), ( + f"After save/load: Expected {expected}, got {reloaded_value[0]}" + ) + finally: + os.unlink(tmp_path) + + +def test_list_item_assignment() -> None: + """Test that list item assignment works for encoder blocks.""" + # Create network A with encoder block containing value 10.0. + net_a = net_pb2.Net() + net_a.format.weights_encoding = net_pb2.Format.LINEAR16 + encoder_a = net_a.weights.encoder.add() + encoder_a.mha.q_w.min_val = 10.0 + encoder_a.mha.q_w.max_val = 10.0 + encoder_a.mha.q_w.params = np.array([32767], dtype=np.uint16).tobytes() + + # Create network B with encoder block containing value 20.0. + net_b = net_pb2.Net() + net_b.format.weights_encoding = net_pb2.Format.LINEAR16 + encoder_b = net_b.weights.encoder.add() + encoder_b.mha.q_w.min_val = 20.0 + encoder_b.mha.q_w.max_val = 20.0 + encoder_b.mha.q_w.params = np.array([32767], dtype=np.uint16).tobytes() + + # Wrap networks. + wrapper_a = NetWrapper(net_a) + wrapper_b = NetWrapper(net_b) + + # Assign encoder[0] from B to A. + wrapper_a.weights.encoder[0] = wrapper_b.weights.encoder[0] + + # Verify in-memory: A's encoder[0] should now have value 20.0. + result_value = wrapper_a.weights.encoder[0].mha.q_w.value + assert np.isclose(result_value[0], 20.0, rtol=1e-4), ( + f"Expected 20.0 (B's value), got {result_value[0]}" + ) + + # Verify persistence. + with tempfile.NamedTemporaryFile(suffix=".pb.gz", delete=False) as tmp: + tmp_path = tmp.name + try: + save_weights(wrapper_a, tmp_path) + reloaded = load_weights(tmp_path) + reloaded_value = reloaded.weights.encoder[0].mha.q_w.value + assert np.isclose(reloaded_value[0], 20.0, rtol=1e-4), ( + f"After save/load: Expected 20.0, got {reloaded_value[0]}" + ) + finally: + os.unlink(tmp_path) diff --git a/src/lczero_training/tools/__init__.py b/src/lczero_training/tools/__init__.py new file mode 100644 index 00000000..12c33a77 --- /dev/null +++ b/src/lczero_training/tools/__init__.py @@ -0,0 +1,5 @@ +"""Pure Python tools for weight manipulation.""" + +from .weights_tool import load_weights, save_weights + +__all__ = ["load_weights", "save_weights"] diff --git a/src/lczero_training/tools/weight_codecs.py b/src/lczero_training/tools/weight_codecs.py new file mode 100644 index 00000000..8e84f56e --- /dev/null +++ b/src/lczero_training/tools/weight_codecs.py @@ -0,0 +1,96 @@ +"""Encoding and decoding logic for Lc0 weight formats.""" + +import numpy as np + +from proto import net_pb2 + + +def decode_linear16( + params: bytes, min_val: float, max_val: float, shape: tuple[int, ...] +) -> np.ndarray: + """Decode LINEAR16 format to float32 array.""" + raw = np.frombuffer(params, dtype=np.uint16) + norm = raw.astype(np.float32) / 65535.0 + values = norm * max_val + (1.0 - norm) * min_val + return values.reshape(shape[::-1]).transpose() + + +def encode_linear16(arr: np.ndarray) -> tuple[bytes, float, float]: + """Encode float32 array to LINEAR16 format.""" + flat = arr.T.flatten().astype(np.float32) + min_val = float(flat.min()) + max_val = float(flat.max()) + rng = max_val - min_val + + if rng < 1e-8: + norm = np.full_like(flat, 0.5) + else: + norm = (flat - min_val) / rng + + quant = np.round(norm * 65535.0).astype(np.uint16) + return quant.tobytes(), min_val, max_val + + +def decode_float16(params: bytes, shape: tuple[int, ...]) -> np.ndarray: + """Decode FLOAT16 format to float32 array.""" + raw = np.frombuffer(params, dtype=np.float16) + values = raw.astype(np.float32) + return values.reshape(shape[::-1]).transpose() + + +def encode_float16(arr: np.ndarray) -> tuple[bytes, float, float]: + """Encode float32 array to FLOAT16 format.""" + flat = arr.T.flatten().astype(np.float16) + return flat.tobytes(), 0.0, 0.0 + + +def decode_bfloat16(params: bytes, shape: tuple[int, ...]) -> np.ndarray: + """Decode BFLOAT16 format to float32 array via bit manipulation.""" + raw_u16 = np.frombuffer(params, dtype=np.uint16) + raw_u32 = raw_u16.astype(np.uint32) << 16 + values = raw_u32.view(np.float32) + return values.reshape(shape[::-1]).transpose() + + +def encode_bfloat16(arr: np.ndarray) -> tuple[bytes, float, float]: + """Encode float32 array to BFLOAT16 format via bit manipulation.""" + flat = arr.T.flatten().astype(np.float32) + u32 = flat.view(np.uint32) + u16 = (u32 >> 16).astype(np.uint16) + return u16.tobytes(), 0.0, 0.0 + + +def decode_layer( + layer: net_pb2.Weights.Layer, fallback_encoding: int +) -> np.ndarray: + """Decode a Layer protobuf to float32 NumPy array.""" + encoding = layer.encoding if layer.encoding else fallback_encoding + + if not layer.dims: + size = len(layer.params) // 2 + shape: tuple[int, ...] = (size,) + else: + shape = tuple(layer.dims) + + if encoding == net_pb2.Weights.Layer.LINEAR16: + return decode_linear16( + layer.params, layer.min_val, layer.max_val, shape + ) + elif encoding == net_pb2.Weights.Layer.FLOAT16: + return decode_float16(layer.params, shape) + elif encoding == net_pb2.Weights.Layer.BFLOAT16: + return decode_bfloat16(layer.params, shape) + else: + raise ValueError(f"Unknown encoding: {encoding}") + + +def encode_layer(arr: np.ndarray, encoding: int) -> tuple[bytes, float, float]: + """Encode a float32 NumPy array to Layer format.""" + if encoding == net_pb2.Weights.Layer.LINEAR16: + return encode_linear16(arr) + elif encoding == net_pb2.Weights.Layer.FLOAT16: + return encode_float16(arr) + elif encoding == net_pb2.Weights.Layer.BFLOAT16: + return encode_bfloat16(arr) + else: + raise ValueError(f"Unknown encoding: {encoding}") diff --git a/src/lczero_training/tools/weight_wrappers.py b/src/lczero_training/tools/weight_wrappers.py new file mode 100644 index 00000000..991ec030 --- /dev/null +++ b/src/lczero_training/tools/weight_wrappers.py @@ -0,0 +1,384 @@ +"""Wrapper classes for pythonic access to Lc0 weight protobufs.""" + +import gzip +from typing import Any, Iterator + +import numpy as np +from google.protobuf.message import Message + +from proto import net_pb2 + +from . import weight_codecs + + +class LayerWrapper: + """Wraps a net_pb2.Weights.Layer with lazy float32 decoding.""" + + __slots__ = ("_proto", "_fallback_encoding", "_cached_array", "_modified") + + def __init__( + self, proto: net_pb2.Weights.Layer, fallback_encoding: int + ) -> None: + object.__setattr__(self, "_proto", proto) + object.__setattr__(self, "_fallback_encoding", fallback_encoding) + object.__setattr__(self, "_cached_array", None) + object.__setattr__(self, "_modified", False) + + _proto: net_pb2.Weights.Layer + _fallback_encoding: int + _cached_array: np.ndarray | None + _modified: bool + + @property + def value(self) -> np.ndarray: + """Decode to float32 on first access.""" + if self._cached_array is None: + decoded = weight_codecs.decode_layer( + self._proto, self._fallback_encoding + ) + object.__setattr__(self, "_cached_array", decoded) + assert self._cached_array is not None + return self._cached_array + + @value.setter + def value(self, arr: np.ndarray) -> None: + """Set new array value, mark as modified.""" + object.__setattr__(self, "_cached_array", arr.astype(np.float32)) + object.__setattr__(self, "_modified", True) + + def commit(self, encoding: int) -> None: + """Re-encode array to proto if modified.""" + if self._modified and self._cached_array is not None: + params, min_val, max_val = weight_codecs.encode_layer( + self._cached_array, encoding + ) + self._proto.params = params + self._proto.min_val = min_val + self._proto.max_val = max_val + self._proto.encoding = encoding # type: ignore[assignment] + del self._proto.dims[:] + self._proto.dims.extend(self._cached_array.shape) + object.__setattr__(self, "_modified", False) + + def __add__(self, other: "LayerWrapper") -> "LayerWrapper": + if not isinstance(other, LayerWrapper): + raise TypeError("Can only add LayerWrapper to LayerWrapper") + result = LayerWrapper(net_pb2.Weights.Layer(), self._fallback_encoding) + result.value = self.value + other.value + return result + + def __sub__(self, other: "LayerWrapper") -> "LayerWrapper": + if not isinstance(other, LayerWrapper): + raise TypeError("Can only subtract LayerWrapper from LayerWrapper") + result = LayerWrapper(net_pb2.Weights.Layer(), self._fallback_encoding) + result.value = self.value - other.value + return result + + def __mul__(self, scalar: float) -> "LayerWrapper": + if not isinstance(scalar, (int, float)): + raise TypeError("Can only multiply LayerWrapper by scalar") + result = LayerWrapper(net_pb2.Weights.Layer(), self._fallback_encoding) + result.value = self.value * scalar + return result + + def __rmul__(self, scalar: float) -> "LayerWrapper": + return self.__mul__(scalar) + + +class ListWrapper: + """Wraps protobuf repeated fields.""" + + __slots__ = ("_proto_list", "_parent", "_item_cache") + + def __init__(self, proto_list: Any, parent: "NetWrapper") -> None: + self._proto_list = proto_list + self._parent = parent + self._item_cache: dict[int, Any] = {} + + _proto_list: Any + _parent: "NetWrapper" + _item_cache: dict[int, Any] + + def __len__(self) -> int: + return len(self._proto_list) + + def __getitem__(self, idx: int) -> Any: + if idx not in self._item_cache: + item_proto = self._proto_list[idx] + self._item_cache[idx] = self._parent._wrap_field(item_proto) + return self._item_cache[idx] + + def __setitem__(self, idx: int, value: Any) -> None: + """Support item assignment for list elements.""" + if isinstance(value, NetWrapper): + dest_proto = self._proto_list[idx] + dest_proto.CopyFrom(value._proto) + # Create new wrapper for destination proto to maintain cache consistency. + self._item_cache[idx] = NetWrapper( + dest_proto, self._parent._fallback_encoding + ) + elif isinstance(value, LayerWrapper): + dest_proto = self._proto_list[idx] + dest_proto.CopyFrom(value._proto) + # Create new wrapper for destination proto to maintain cache consistency. + self._item_cache[idx] = LayerWrapper( + dest_proto, self._parent._fallback_encoding + ) + else: + self._proto_list[idx] = value + if idx in self._item_cache: + del self._item_cache[idx] + + def __iter__(self) -> Iterator[Any]: + for i in range(len(self)): + yield self[i] + + +class NetWrapper: + """Wraps net_pb2.Net or nested Message types.""" + + __slots__ = ("_proto", "_fallback_encoding", "_attr_cache") + + def __init__( + self, proto_msg: Message, fallback_encoding: int | None = None + ) -> None: + object.__setattr__(self, "_proto", proto_msg) + object.__setattr__( + self, + "_fallback_encoding", + fallback_encoding or self._detect_encoding(), + ) + object.__setattr__(self, "_attr_cache", {}) + + _proto: Message + _fallback_encoding: int + _attr_cache: dict[str, Any] + + def _detect_encoding(self) -> int: + """Extract encoding from net.format.weights_encoding.""" + if hasattr(self._proto, "format") and self._proto.HasField("format"): + return self._proto.format.weights_encoding + return net_pb2.Weights.Layer.LINEAR16 + + def __getattr__(self, name: str) -> Any: + if name.startswith("_"): + return object.__getattribute__(self, name) + + if name in self._attr_cache: + return self._attr_cache[name] + + if not hasattr(self._proto, name): + raise AttributeError( + f"{type(self._proto).__name__} has no field '{name}'" + ) + + value = getattr(self._proto, name) + wrapped = self._wrap_field(value) + self._attr_cache[name] = wrapped + return wrapped + + def _wrap_field(self, value: Any) -> Any: + """Determine wrapper type based on proto field.""" + if isinstance(value, net_pb2.Weights.Layer): + return LayerWrapper(value, self._fallback_encoding) + elif isinstance(value, Message): + return NetWrapper(value, self._fallback_encoding) + elif hasattr(value, "__len__") and not isinstance(value, (str, bytes)): + return ListWrapper(value, self) + else: + return value + + def __setattr__(self, name: str, value: Any) -> None: + if name.startswith("_"): + object.__setattr__(self, name, value) + return + + if isinstance(value, NetWrapper): + dest_proto = getattr(self._proto, name) + dest_proto.CopyFrom(value._proto) + # Create new wrapper for destination proto to maintain cache consistency. + self._attr_cache[name] = NetWrapper( + dest_proto, self._fallback_encoding + ) + elif isinstance(value, LayerWrapper): + dest_proto = getattr(self._proto, name) + dest_proto.CopyFrom(value._proto) + # Create new wrapper for destination proto to maintain cache consistency. + self._attr_cache[name] = LayerWrapper( + dest_proto, self._fallback_encoding + ) + else: + setattr(self._proto, name, value) + + def save(self, path: str, encoding: int | None = None) -> None: + """Save to file, committing all modified layers.""" + if encoding is None: + encoding = net_pb2.Weights.Layer.FLOAT16 + + self._commit_all(encoding) + + serialized = self._proto.SerializePartialToString() + if path.endswith(".gz"): + with gzip.open(path, "wb") as f: + f.write(serialized) + else: + with open(path, "wb") as f: + f.write(serialized) + + def _commit_all(self, encoding: int) -> None: + """Recursively commit all modified LayerWrappers.""" + for cached_value in self._attr_cache.values(): + if isinstance(cached_value, LayerWrapper): + cached_value.commit(encoding) + elif isinstance(cached_value, NetWrapper): + cached_value._commit_all(encoding) + elif isinstance(cached_value, ListWrapper): + for item in cached_value: + if isinstance(item, NetWrapper): + item._commit_all(encoding) + elif isinstance(item, LayerWrapper): + item.commit(encoding) + + def __add__(self, other: "NetWrapper") -> "NetWrapper": + """Element-wise addition of two networks.""" + if not isinstance(other, NetWrapper): + raise TypeError("Can only add NetWrapper to NetWrapper") + + result_proto = type(self._proto)() + result_proto.CopyFrom(self._proto) + result = NetWrapper(result_proto, self._fallback_encoding) + result._add_weights(self, other) + return result + + def __sub__(self, other: "NetWrapper") -> "NetWrapper": + """Element-wise subtraction of two networks.""" + if not isinstance(other, NetWrapper): + raise TypeError("Can only subtract NetWrapper from NetWrapper") + + result_proto = type(self._proto)() + result_proto.CopyFrom(self._proto) + result = NetWrapper(result_proto, self._fallback_encoding) + result._sub_weights(self, other) + return result + + def __mul__(self, scalar: float) -> "NetWrapper": + """Scalar multiplication.""" + if not isinstance(scalar, (int, float)): + raise TypeError("Can only multiply NetWrapper by scalar") + + result_proto = type(self._proto)() + result_proto.CopyFrom(self._proto) + result = NetWrapper(result_proto, self._fallback_encoding) + result._mul_weights(self, scalar) + return result + + def __rmul__(self, scalar: float) -> "NetWrapper": + return self.__mul__(scalar) + + def _add_weights(self, lhs: "NetWrapper", rhs: "NetWrapper") -> None: + """Recursively add weights from lhs and rhs into self.""" + for field_desc in lhs._proto.DESCRIPTOR.fields: + field_name = field_desc.name + if not hasattr(lhs._proto, field_name): + continue + + # Skip optional fields that are not set in both inputs. + if not field_desc.is_required and not field_desc.is_repeated: + lhs_has = lhs._proto.HasField(field_name) + rhs_has = rhs._proto.HasField(field_name) + if not lhs_has or not rhs_has: + continue + + lhs_val = getattr(lhs, field_name) + rhs_val = getattr(rhs, field_name) + + if isinstance(lhs_val, LayerWrapper): + self_layer = getattr(self, field_name) + self_layer.value = lhs_val.value + rhs_val.value + elif isinstance(lhs_val, NetWrapper): + self_wrapper = getattr(self, field_name) + self_wrapper._add_weights(lhs_val, rhs_val) + elif isinstance(lhs_val, ListWrapper): + self_list = getattr(self, field_name) + # Only process indices that exist in both lists. + min_len = min(len(lhs_val), len(rhs_val)) + for i in range(min_len): + if isinstance(lhs_val[i], (NetWrapper, LayerWrapper)): + if isinstance(lhs_val[i], LayerWrapper): + self_list[i].value = ( + lhs_val[i].value + rhs_val[i].value + ) + else: + self_list[i]._add_weights(lhs_val[i], rhs_val[i]) + # Truncate to min_len to avoid incomplete entries. + del self_list._proto_list[min_len:] + + def _sub_weights(self, lhs: "NetWrapper", rhs: "NetWrapper") -> None: + """Recursively subtract weights rhs from lhs into self.""" + for field_desc in lhs._proto.DESCRIPTOR.fields: + field_name = field_desc.name + if not hasattr(lhs._proto, field_name): + continue + + # Skip optional fields that are not set in both inputs. + if not field_desc.is_required and not field_desc.is_repeated: + lhs_has = lhs._proto.HasField(field_name) + rhs_has = rhs._proto.HasField(field_name) + if not lhs_has or not rhs_has: + continue + + lhs_val = getattr(lhs, field_name) + rhs_val = getattr(rhs, field_name) + + if isinstance(lhs_val, LayerWrapper): + self_layer = getattr(self, field_name) + self_layer.value = lhs_val.value - rhs_val.value + elif isinstance(lhs_val, NetWrapper): + self_wrapper = getattr(self, field_name) + self_wrapper._sub_weights(lhs_val, rhs_val) + elif isinstance(lhs_val, ListWrapper): + self_list = getattr(self, field_name) + # Only process indices that exist in both lists. + min_len = min(len(lhs_val), len(rhs_val)) + for i in range(min_len): + if isinstance(lhs_val[i], (NetWrapper, LayerWrapper)): + if isinstance(lhs_val[i], LayerWrapper): + self_list[i].value = ( + lhs_val[i].value - rhs_val[i].value + ) + else: + self_list[i]._sub_weights(lhs_val[i], rhs_val[i]) + # Truncate to min_len to avoid incomplete entries. + del self_list._proto_list[min_len:] + + def _mul_weights(self, source: "NetWrapper", scalar: float) -> None: + """Recursively multiply all weights by scalar.""" + for field_desc in source._proto.DESCRIPTOR.fields: + field_name = field_desc.name + if not hasattr(source._proto, field_name): + continue + + # Skip unset optional fields to avoid creating them in output. + if ( + not field_desc.is_required + and not field_desc.is_repeated + and not source._proto.HasField(field_name) + ): + continue + + src_val = getattr(source, field_name) + + if isinstance(src_val, LayerWrapper): + self_layer = getattr(self, field_name) + self_layer.value = src_val.value * scalar + elif isinstance(src_val, NetWrapper): + self_wrapper = getattr(self, field_name) + self_wrapper._mul_weights(src_val, scalar) + elif isinstance(src_val, ListWrapper): + self_list = getattr(self, field_name) + for i in range(len(src_val)): + if isinstance(src_val[i], (NetWrapper, LayerWrapper)): + if isinstance(src_val[i], LayerWrapper): + self_list[i].value = src_val[i].value * scalar + else: + self_list[i]._mul_weights(src_val[i], scalar) diff --git a/src/lczero_training/tools/weights_tool.py b/src/lczero_training/tools/weights_tool.py new file mode 100644 index 00000000..3ee90562 --- /dev/null +++ b/src/lczero_training/tools/weights_tool.py @@ -0,0 +1,34 @@ +"""Main API for loading and saving Lc0 weight files.""" + +import gzip + +from proto import net_pb2 + +from .weight_wrappers import NetWrapper + + +def load_weights(path: str) -> NetWrapper: + """Load Lc0 weights file (.pb or .pb.gz).""" + if path.endswith(".gz"): + with gzip.open(path, "rb") as f: + contents = f.read() + else: + with open(path, "rb") as f: + contents = f.read() + + net = net_pb2.Net() + net.ParseFromString(contents) + return NetWrapper(net) + + +def save_weights( + wrapper: NetWrapper, path: str, encoding: str = "FLOAT16" +) -> None: + """Save weights to file.""" + encoding_map = { + "LINEAR16": net_pb2.Weights.Layer.LINEAR16, + "FLOAT16": net_pb2.Weights.Layer.FLOAT16, + "BFLOAT16": net_pb2.Weights.Layer.BFLOAT16, + } + enc_value = encoding_map[encoding.upper()] + wrapper.save(path, enc_value) diff --git a/src/lczero_training/training/__init__.py b/src/lczero_training/training/__init__.py new file mode 100644 index 00000000..8c7eb32c --- /dev/null +++ b/src/lczero_training/training/__init__.py @@ -0,0 +1 @@ +"""Training package for Leela Chess Zero.""" diff --git a/src/lczero_training/training/backfill_metrics.py b/src/lczero_training/training/backfill_metrics.py new file mode 100644 index 00000000..79d9b536 --- /dev/null +++ b/src/lczero_training/training/backfill_metrics.py @@ -0,0 +1,148 @@ +"""Backfill metrics for existing checkpoints.""" + +import logging +from typing import Any + +from flax import nnx +from google.protobuf import text_format + +from lczero_training.daemon.metrics import ( + evaluate_batch, + load_batch_from_npz, +) +from lczero_training.model.loss_function import LczeroLoss +from lczero_training.model.model import LczeroModel +from lczero_training.training.migrate_checkpoint import ( + Migration, + get_checkpoint_steps, + load_checkpoint, + load_migration_rules, +) +from lczero_training.training.state import TrainingState +from lczero_training.training.tensorboard import TensorboardLogger +from proto.root_config_pb2 import RootConfig + +logger = logging.getLogger(__name__) + + +def _load_config(config_path: str) -> RootConfig: + """Load RootConfig from textproto file.""" + config = RootConfig() + with open(config_path, "r") as f: + text_format.Parse(f.read(), config) + return config + + +def _validate_and_get_metrics( + root_config: RootConfig, metric_names: list[str] +) -> dict[str, Any]: + """Validate metrics exist and are NPZ type.""" + if not root_config.HasField("metrics"): + raise ValueError("No metrics configuration found in root config") + + metric_configs = { + mc.name: mc + for mc in root_config.metrics.metric + if mc.name in metric_names and mc.HasField("npz_filename") + } + + non_npz = [ + mc.name + for mc in root_config.metrics.metric + if mc.name in metric_names and not mc.HasField("npz_filename") + ] + if non_npz: + raise ValueError(f"Non-NPZ metrics: {', '.join(non_npz)}") + + missing = set(metric_names) - set(metric_configs.keys()) + if missing: + raise ValueError(f"Metrics not found: {', '.join(sorted(missing))}") + + return metric_configs + + +def _load_and_migrate_checkpoint( + checkpoint_path: str, + step: int, + template: TrainingState, + rules: list[tuple], +) -> TrainingState: + """Load checkpoint and apply migration if rules provided.""" + state, _ = load_checkpoint(checkpoint_path, step) + return Migration(state, template).run(rules) + + +def backfill_metrics( + config_path: str, + metric_names: list[str], + min_step: int | None = None, + max_step: int | None = None, + migration_config_path: str | None = None, +) -> None: + """Backfill metrics for existing checkpoints. + + Args: + config_path: Path to the RootConfig textproto file. + metric_names: Names of metrics to backfill (must be NPZ metrics). + min_step: Minimum checkpoint step (inclusive), or None. + max_step: Maximum checkpoint step (inclusive), or None. + migration_config_path: Path to CheckpointMigrationConfig file, or None. + + Raises: + ValueError: If any metric is not an NPZ metric or doesn't exist. + """ + config = _load_config(config_path) + metric_configs = _validate_and_get_metrics(config, metric_names) + + # Load batches and create loggers. + batches = { + name: load_batch_from_npz(mc.npz_filename) + for name, mc in metric_configs.items() + } + loggers = { + name: TensorboardLogger(f"{config.metrics.tensorboard_path}/{name}") + for name in metric_configs + } + + # Initialize model components. + loss_fn = LczeroLoss(config=config.training.losses) + graphdef, _ = nnx.split( + LczeroModel(config=config.model, rngs=nnx.Rngs(params=42)) + ) + + # Prepare migration if needed. + rules = load_migration_rules(migration_config_path) + template = TrainingState.new_from_config(config.model, config.training) + + # Get and process checkpoints. + steps = get_checkpoint_steps( + config.training.checkpoint.path, min_step, max_step + ) + logger.info(f"Processing {len(steps)} checkpoints") + + if not steps: + logger.warning("No checkpoints found in range") + return + + try: + for step in steps: + logger.info(f"Step {step}") + state = _load_and_migrate_checkpoint( + config.training.checkpoint.path, step, template, rules + ) + + for name, mc in metric_configs.items(): + metrics = evaluate_batch( + batches[name], + state.jit_state, + graphdef, + loss_fn, + mc.use_swa_model, + ) + loggers[name].log(step, metrics) + logger.info(f" {name}: loss={metrics['loss']:.6f}") + finally: + for tb_logger in loggers.values(): + tb_logger.close() + + logger.info("Backfill complete") diff --git a/src/lczero_training/training/dataloader_probe.py b/src/lczero_training/training/dataloader_probe.py new file mode 100644 index 00000000..b7645cc8 --- /dev/null +++ b/src/lczero_training/training/dataloader_probe.py @@ -0,0 +1,99 @@ +"""Utilities for exercising the training data loader.""" + +import logging +import time +from contextlib import suppress +from pathlib import Path +from typing import Optional + +import numpy as np +from google.protobuf import text_format + +from lczero_training.dataloader import DataLoader, make_dataloader +from proto.root_config_pb2 import RootConfig + +logger = logging.getLogger(__name__) + + +def _stop_loader(loader: DataLoader) -> None: + with suppress(Exception): + loader.stop() + + +def _store_batches(path: str, batches: list) -> None: + output = Path(path) + if output.parent: + output.parent.mkdir(parents=True, exist_ok=True) + logger.info("Writing %d batches to %s", len(batches), output) + container = np.empty(len(batches), dtype=object) + container[:] = batches + np.savez(output, batches=container) + + +def probe_dataloader( + config_filename: str, num_batches: int, npz_output: Optional[str] = None +) -> None: + """Measure latency and throughput for the configured data loader. + + Args: + config_filename: Path to the root configuration proto file. + num_batches: Total number of batches to fetch from the loader. + npz_output: Optional path to store fetched batches as an .npz archive. + """ + + if num_batches < 1: + raise ValueError("num_batches must be at least 1") + + config = RootConfig() + logger.info("Reading configuration from proto file") + with open(config_filename, "r") as config_file: + text_format.Parse(config_file.read(), config) + + logger.info("Creating data loader") + loader = make_dataloader(config.data_loader) + + collected_batches: list = [] + collect_enabled = npz_output is not None + first_batch_time = 0.0 + remaining_batches = num_batches - 1 + try: + logger.info("Fetching first batch") + start_time = time.perf_counter() + first_batch = loader.get_next() + if collect_enabled: + collected_batches.append(first_batch) + first_batch_time = time.perf_counter() - start_time + logger.info("Time to first batch: %.3f seconds", first_batch_time) + + if remaining_batches <= 0: + logger.info("Only fetched first batch; skipping throughput") + return + + logger.info( + "Fetching %d additional batches for throughput measurement", + remaining_batches, + ) + throughput_start = time.perf_counter() + for _ in range(remaining_batches): + batch = loader.get_next() + if collect_enabled: + collected_batches.append(batch) + throughput_duration = time.perf_counter() - throughput_start + + if throughput_duration <= 0: + logger.warning("Measured non-positive duration; skipping rate") + return + + batches_per_second = remaining_batches / throughput_duration + logger.info( + "Throughput excluding first batch: %.2f batches/second", + batches_per_second, + ) + logger.info( + "Total time excluding first batch: %.3f seconds", + throughput_duration, + ) + finally: + if collect_enabled and npz_output is not None and collected_batches: + _store_batches(npz_output, collected_batches) + _stop_loader(loader) diff --git a/src/lczero_training/training/describe.py b/src/lczero_training/training/describe.py new file mode 100644 index 00000000..4737a3d3 --- /dev/null +++ b/src/lczero_training/training/describe.py @@ -0,0 +1,77 @@ +import logging +import sys +from pathlib import PurePosixPath + +import jax +import jax.numpy as jnp +import orbax.checkpoint as ocp +from flax import nnx +from google.protobuf import text_format + +from lczero_training.training.state import TrainingState +from proto.root_config_pb2 import RootConfig + +logger = logging.getLogger(__name__) + + +def describe( + config_filename: str, + shapes: bool = False, + values: bool = False, + weight_paths: bool = False, +) -> None: + config = RootConfig() + logger.info("Reading configuration from proto file") + with open(config_filename, "r") as f: + text_format.Parse(f.read(), config) + + if config.training.checkpoint.path is None: + logger.error("Checkpoint path must be set in the configuration.") + sys.exit(1) + + checkpoint_mgr = ocp.CheckpointManager( + config.training.checkpoint.path, + options=ocp.CheckpointManagerOptions( + create=True, + ), + ) + + logger.info("Creating state from configuration") + empty_state = TrainingState.new_from_config( + model_config=config.model, + training_config=config.training, + ) + logger.info("Restoring checkpoint") + training_state = checkpoint_mgr.restore( + None, args=ocp.args.PyTreeRestore(empty_state) + ) + logger.info("Restored checkpoint") + + assert isinstance(training_state, TrainingState) + + if values: + logger.info("Dumping training state values") + print("Training state:") + print(training_state) + + if shapes: + logger.info("Extracting training state shapes") + shapes = jax.tree.map(jnp.shape, training_state) + print("Training state shapes:") + print(shapes) + + if weight_paths: + paths = [] + + def _collect(path: tuple[object, ...], _: nnx.Variable) -> bool: + paths.append(str(PurePosixPath(*map(str, path)))) + return False + + nnx.map_state(_collect, training_state.jit_state.model_state) + for p in sorted( + paths, + key=lambda p: tuple( + int(c) if c.isdigit() else c for c in p.split("/") + ), + ): + print(p) diff --git a/src/lczero_training/training/eval.py b/src/lczero_training/training/eval.py new file mode 100644 index 00000000..b43acd47 --- /dev/null +++ b/src/lczero_training/training/eval.py @@ -0,0 +1,606 @@ +# Description: Evaluation script for comparing model outputs and calculating losses. +# +# This script provides functionalities to evaluate a trained model by processing +# data samples, calculating losses, and comparing outputs against an ONNX model. +# It supports dumping tensors and results to various formats for analysis. + +import json +import logging +import math +import shelve +import sys +from dataclasses import dataclass +from datetime import datetime +from typing import ( + Any, + Callable, + Dict, + Generator, + Optional, + Sequence, + TextIO, + Tuple, + cast, +) + +import jax +import jax.numpy as jnp +import numpy as np +import orbax.checkpoint as ocp +from flax import nnx +from google.protobuf import text_format + +from lczero_training.dataloader import ( + DataLoader, + make_dataloader, +) +from lczero_training.model.loss_function import LczeroLoss +from lczero_training.model.model import LczeroModel, ModelPrediction +from lczero_training.training.state import TrainingSample, TrainingState +from proto import data_loader_config_pb2 +from proto.root_config_pb2 import RootConfig + +logger = logging.getLogger(__name__) + +RELATIVE_DIFFERENCE_EPSILON = 1e-4 +HEADS = ("wdl", "policy", "movesleft") + + +@dataclass +class DiffRecord: + """Stores information about the difference between JAX and ONNX outputs.""" + + batch: int + sample: int + index: Tuple[int, ...] + diff: float + jax_value: float + onnx_value: float + + +def _tensor_to_list(obj: Any) -> Any: + """Recursively converts JAX/Numpy arrays to Python lists for serialization.""" + if hasattr(obj, "tolist"): + return obj.tolist() + if isinstance(obj, dict): + return {key: _tensor_to_list(value) for key, value in obj.items()} + return obj + + +# --- Diff statistics helpers --- + + +def _bin_counts(values: np.ndarray) -> Dict[str, Any]: + """Counts values into bins based on powers of 2.""" + flat = np.asarray(values).ravel() + zero_count = int(np.count_nonzero(flat == 0.0)) + non_zero = flat[flat != 0.0] + + bins: Dict[int, int] = {} + if non_zero.size > 0: + exponents = np.floor(np.log2(non_zero)).astype(int) + unique, counts = np.unique(exponents, return_counts=True) + bins = {int(exp): int(count) for exp, count in zip(unique, counts)} + + return {"zero": zero_count, "bins": bins} + + +def _format_bound(value: float) -> str: + """Formats a float for display in statistics.""" + if value >= 1 and math.isclose(value, round(value)): + return str(int(round(value))) + return f"{value:.6g}" + + +def _format_stats(stats: Dict[str, Any]) -> str: + """Formats bin statistics into a readable string.""" + lines = [f" zero={stats['zero']}"] + for exponent in sorted(stats["bins"].keys(), reverse=True): + lower = 2**exponent + upper = 2 ** (exponent + 1) + lines.append( + f" [{_format_bound(lower)}; {_format_bound(upper)})=" + f"{stats['bins'][exponent]}" + ) + return "\n".join(lines) + + +def _collect_diff_statistics( + jax_output: np.ndarray, onnx_output: np.ndarray +) -> Tuple[np.ndarray, Dict[str, Any], Optional[Dict[str, Any]]]: + """Collects absolute and relative difference statistics.""" + abs_diff = np.abs(jax_output - onnx_output) + abs_stats = _bin_counts(abs_diff) + + mask = np.abs(jax_output) >= RELATIVE_DIFFERENCE_EPSILON + if not np.any(mask): + return abs_diff, abs_stats, None + + rel_diff = np.zeros_like(abs_diff, dtype=np.float64) + np.divide(abs_diff, np.abs(jax_output), out=rel_diff, where=mask) + rel_stats = _bin_counts(rel_diff[mask]) + return abs_diff, abs_stats, rel_stats + + +class Dumper: + """Handles dumping of evaluation artifacts.""" + + def __init__( + self, + to_stdout: bool, + to_file: Optional[str], + to_shelve: Optional[str], + to_json: Optional[str], + ): + self.to_stdout = to_stdout + self.shelve_path = to_shelve + self.json_path = to_json + self.file_handle: Optional[TextIO] = ( + open(to_file, "w") if to_file else None + ) + + def dump_tensors(self, tensors: dict, prefix: str) -> None: + """Dumps tensors to stdout or a text file.""" + if not self.to_stdout and not self.file_handle: + return + + lines = [f"# === {prefix} TENSORS ==="] + for name, tensor in tensors.items(): + lines.append(f"{name} = {str(_tensor_to_list(tensor))}") + lines.append("") + output_text = "\n".join(lines) + + if self.to_stdout: + print(output_text) + if self.file_handle: + self.file_handle.write(output_text) + self.file_handle.flush() + + def dump_structured(self, batch: dict, outputs: dict, losses: dict) -> None: + """Dumps results to structured formats like JSON or shelve.""" + if not self.shelve_path and not self.json_path: + return + + all_data = { + **_tensor_to_list(batch), + **_tensor_to_list(outputs), + **_tensor_to_list(losses), + } + key = f"sample-{datetime.now().strftime('%Y%m%d-%H%M%S')}" + + if self.shelve_path: + self._dump_to_shelve(key, all_data) + if self.json_path: + self._dump_to_json(key, all_data) + + def _dump_to_shelve(self, key: str, data: dict) -> None: + assert self.shelve_path is not None + with shelve.open(self.shelve_path) as db: + db[key] = data + logger.info("Dumped data to shelve with key: %s", key) + + def _dump_to_json(self, key: str, data: dict) -> None: + assert self.json_path is not None + try: + with open(self.json_path, "r") as f: + json_data = json.load(f) + except (FileNotFoundError, json.JSONDecodeError): + json_data = {} + json_data[key] = data + with open(self.json_path, "w") as f: + json.dump(json_data, f, indent=2) + logger.info("Dumped data to JSON with key: %s", key) + + def close(self) -> None: + if self.file_handle: + self.file_handle.close() + + +class OnnxComparator: + """Handles comparison of JAX model outputs with ONNX model outputs.""" + + def __init__(self, onnx_model_path: str): + try: + import onnxruntime as ort + except ImportError as exc: + raise RuntimeError( + "onnxruntime is required for ONNX comparison." + ) from exc + + self.session = ort.InferenceSession(onnx_model_path) + inputs = self.session.get_inputs() + if not inputs: + raise ValueError("ONNX model must define at least one input.") + self.input_name = inputs[0].name + logger.info("Loaded ONNX model for comparison from %s", onnx_model_path) + + self.head_mapping_logged = False + self.worst_records: Dict[str, Optional[DiffRecord]] = { + head: None for head in HEADS + } + self.onnx_outputs: Dict[str, jax.Array] = {} + + def compare( + self, + jax_outputs: Dict[str, jax.Array], + onnx_inputs_np: np.ndarray, + sample_index: int, + ) -> None: + """Runs comparison for a single sample.""" + raw_onnx_outputs = self.session.run( + None, {self.input_name: onnx_inputs_np} + ) + if len(raw_onnx_outputs) != 3: + raise ValueError( + "Expected three outputs (wdl, policy, movesleft) from ONNX model." + ) + + jax_outputs_np = {k: np.asarray(v) for k, v in jax_outputs.items()} + aligned_onnx, head_indices = self._align_onnx_outputs( + jax_outputs_np, raw_onnx_outputs + ) + + if not self.head_mapping_logged: + order = ", ".join(f"{h}=output[{head_indices[h]}]" for h in HEADS) + logger.info("Aligned ONNX outputs to heads as: %s", order) + self.head_mapping_logged = True + + self.onnx_outputs = { + "onnx_value_pred": jnp.asarray(aligned_onnx["wdl"]), + "onnx_policy_pred": jnp.asarray(aligned_onnx["policy"]), + "onnx_movesleft_pred": jnp.asarray(aligned_onnx["movesleft"]), + } + + for head in HEADS: + record = self._log_diff_stats( + head, jax_outputs_np[head], aligned_onnx[head], sample_index + ) + current = self.worst_records[head] + if current is None or record.diff > current.diff: + self.worst_records[head] = record + + def log_summary(self) -> None: + """Logs the worst difference found for each head.""" + for head in HEADS: + record = self.worst_records[head] + if record: + logger.info( + "Worst ONNX abs diff for %s head: " + "batch=%d sample=%d index=%s diff=%0.6g " + "jax=%0.6g onnx=%0.6g", + head, + record.batch, + record.sample, + record.index, + record.diff, + record.jax_value, + record.onnx_value, + ) + + def _log_diff_stats( + self, + head: str, + jax_output: np.ndarray, + onnx_output: np.ndarray, + sample_index: int, + ) -> DiffRecord: + if jax_output.shape != onnx_output.shape: + raise ValueError( + f"Shape mismatch for {head} head: " + f"JAX {jax_output.shape} vs ONNX {onnx_output.shape}." + ) + + abs_diff, abs_stats, rel_stats = _collect_diff_statistics( + jax_output, onnx_output + ) + + logger.info( + "Batch %d %s head ONNX abs diff stats:\n%s", + sample_index, + head, + _format_stats(abs_stats), + ) + if rel_stats: + logger.info( + "Batch %d %s head ONNX rel diff stats:\n%s", + sample_index, + head, + _format_stats(rel_stats), + ) + else: + logger.info( + "Batch %d %s head ONNX rel diff stats: skipped (all |jax| < %.1e)", + sample_index, + head, + RELATIVE_DIFFERENCE_EPSILON, + ) + + max_loc = np.unravel_index(int(np.argmax(abs_diff)), abs_diff.shape) + return DiffRecord( + batch=sample_index, + sample=int(max_loc[0]) if max_loc else 0, + index=tuple(int(i) for i in max_loc), + diff=float(abs_diff[max_loc]), + jax_value=float(jax_output[max_loc]), + onnx_value=float(onnx_output[max_loc]), + ) + + def _align_onnx_outputs( + self, + jax_outputs: Dict[str, np.ndarray], + onnx_outputs: Sequence[np.ndarray], + ) -> Tuple[Dict[str, np.ndarray], Dict[str, int]]: + """Matches ONNX outputs to heads regardless of ordering differences.""" + remaining = [(i, np.asarray(o)) for i, o in enumerate(onnx_outputs)] + aligned: Dict[str, np.ndarray] = {} + indices: Dict[str, int] = {} + + def pop_match( + predicate: Callable[[np.ndarray], bool], + ) -> Optional[Tuple[int, np.ndarray]]: + for i, candidate in enumerate(remaining): + if predicate(candidate[1]): + return remaining.pop(i) + return None + + for head in HEADS: + shape = jax_outputs[head].shape + match = pop_match(lambda arr: arr.shape == shape) + if match: + idx, array = match + aligned[head], indices[head] = array, idx + + for head in HEADS: + if head in aligned: + continue + size = jax_outputs[head].size + match = pop_match(lambda arr: arr.size == size) + if match: + idx, array = match + aligned[head], indices[head] = array, idx + + if len(aligned) != len(HEADS): + rem_shapes = [arr.shape for _, arr in remaining] + raise ValueError( + "Could not align ONNX outputs with JAX outputs. " + f"Aligned: {list(aligned.keys())}; Unmatched shapes: {rem_shapes}" + ) + + for head in HEADS: + aligned[head] = self._reshape_output( + aligned[head], jax_outputs[head].shape, head + ) + return aligned, indices + + def _reshape_output( + self, array: np.ndarray, target_shape: Tuple[int, ...], head: str + ) -> np.ndarray: + if array.shape == target_shape: + return array + if array.size != int(np.prod(target_shape)): + raise ValueError( + f"Cannot reshape ONNX output for {head}: source shape " + f"{array.shape}, target shape {target_shape}." + ) + try: + return np.reshape(array, target_shape) + except ValueError as exc: + raise ValueError( + f"Failed to reshape ONNX output for {head} to {target_shape}." + ) from exc + + +class Evaluation: + """Orchestrates the model evaluation process.""" + + def __init__(self, loss_fn: LczeroLoss): + self.loss_fn = loss_fn + + def run( + self, + model: LczeroModel, + datagen: Generator[tuple[np.ndarray, ...], None, None], + num_samples: int, + dumper: Dumper, + onnx_comparator: Optional[OnnxComparator], + softmax_jax_wdl: bool, + ) -> None: + loss_vfn = jax.vmap(self._loss_for_grad, in_axes=(None, 0), out_axes=0) + model_output_vfn = jax.vmap( + self._model_for_output, in_axes=(None, 0), out_axes=0 + ) + + for i in range(num_samples): + logger.info("Processing sample %d/%d", i, num_samples) + self._process_sample( + model, + datagen, + i, + dumper, + onnx_comparator, + loss_vfn, + model_output_vfn, + softmax_jax_wdl, + ) + logger.info("Sample %d complete", i) + + def _process_sample( + self, + model: LczeroModel, + datagen: Generator[tuple[np.ndarray, ...], None, None], + sample_idx: int, + dumper: Dumper, + onnx_comparator: Optional[OnnxComparator], + loss_vfn: Callable[ + [LczeroModel, TrainingSample], + Tuple[jax.Array, Dict[str, jax.Array]], + ], + model_output_vfn: Callable[ + [LczeroModel, jax.Array], + ModelPrediction, + ], + softmax_jax_wdl: bool, + ) -> None: + batch_tuple = next(datagen) + logger.info("Fetched batch from dataloader") + + # DataLoader now returns tuple: (inputs, probabilities, values) + batch = { + "inputs": cast(jax.Array, jnp.asarray(batch_tuple[0])), + "probabilities": cast(jax.Array, jnp.asarray(batch_tuple[1])), + "values": cast(jax.Array, jnp.asarray(batch_tuple[2])), + } + dumper.dump_tensors(batch, "INPUT") + + predictions = model_output_vfn(model, cast(jax.Array, batch["inputs"])) + value_preds = predictions.value + policy_preds = predictions.policy + movesleft_preds = predictions.movesleft + + # Flatten all head outputs for dumping + outputs = {} + for name, pred_tuple in value_preds.items(): + pred = pred_tuple[0] + if softmax_jax_wdl: + pred = jax.nn.softmax(pred, axis=-1) + outputs[f"value_pred/{name}"] = pred + for name, pred in policy_preds.items(): + outputs[f"policy_pred/{name}"] = pred + for name, pred in movesleft_preds.items(): + outputs[f"movesleft_pred/{name}"] = pred + + if onnx_comparator: + # Compare only legacy heads + jax_outputs_for_onnx = { + "wdl": value_preds["winner"][0], + "policy": policy_preds["vanilla"], + "movesleft": movesleft_preds["main"], + } + onnx_inputs_np = np.asarray(batch["inputs"]).copy() + onnx_inputs_np[:, 109, ...] *= 99 + onnx_comparator.compare( + jax_outputs_for_onnx, onnx_inputs_np, sample_idx + ) + outputs.update(onnx_comparator.onnx_outputs) + + dumper.dump_tensors(outputs, "OUTPUT") + + # Convert batch dict to TrainingSample for loss function + batch_sample = TrainingSample( + inputs=batch["inputs"], + probabilities=batch["probabilities"], + values=batch["values"], + ) + per_sample_loss, unweighted_losses = loss_vfn(model, batch_sample) + losses = { + "per_sample_data_loss": per_sample_loss, + "unweighted_losses": unweighted_losses, + } + dumper.dump_tensors(losses, "LOSSES") + dumper.dump_structured(batch, outputs, losses) + + def _loss_for_grad( + self, model_arg: LczeroModel, sample_arg: TrainingSample + ) -> Tuple[jax.Array, Dict[str, jax.Array]]: + return self.loss_fn(model_arg, sample_arg) + + @staticmethod + def _model_for_output( + model_arg: LczeroModel, inputs_arg: jax.Array + ) -> ModelPrediction: + return model_arg(inputs_arg) + + +def from_dataloader( + loader: DataLoader, +) -> Generator[tuple[np.ndarray, ...], None, None]: + """Infinetely yields batches from a DataLoader.""" + while True: + yield loader.get_next() + + +def _load_model_from_checkpoint(config: RootConfig) -> LczeroModel: + """Loads a model from the latest checkpoint.""" + if not config.training.checkpoint.path: + logger.error("Checkpoint path must be set in the configuration.") + sys.exit(1) + + mgr = ocp.CheckpointManager( + config.training.checkpoint.path, + options=ocp.CheckpointManagerOptions(create=True), + ) + state = TrainingState.new_from_config(config.model, config.training) + restored_state = mgr.restore( + mgr.latest_step(), args=ocp.args.PyTreeRestore(state) + ) + logger.info("Restored checkpoint from %s", config.training.checkpoint.path) + + assert isinstance(restored_state, TrainingState) + model_graph, _ = nnx.split( + LczeroModel(config.model, rngs=nnx.Rngs(params=42)) + ) + return nnx.merge(model_graph, restored_state.jit_state.model_state) + + +def _get_dataloader_config( + config: RootConfig, batch_size_override: Optional[int] +) -> data_loader_config_pb2.DataLoaderConfig: + """Gets the dataloader config, overriding batch size if specified.""" + dl_config = config.data_loader + if batch_size_override is None: + return dl_config + + for stage in dl_config.stage: + if stage.HasField("tensor_generator"): + stage.tensor_generator.batch_size = batch_size_override + logger.info("Overriding batch size to %d", batch_size_override) + return dl_config + + raise ValueError( + "tensor_generator stage is required to override batch size" + ) + + +def eval( + config_filename: str, + num_samples: Optional[int] = None, + batch_size_override: Optional[int] = None, + dump_to_stdout: bool = False, + dump_to_file: Optional[str] = None, + dump_to_shelve: Optional[str] = None, + dump_to_json: Optional[str] = None, + onnx_model: Optional[str] = None, + softmax_jax_wdl: bool = True, +) -> None: + """Main function to run the evaluation.""" + jnp.set_printoptions(threshold=sys.maxsize, suppress=False) + + config = RootConfig() + logger.info("Reading configuration from: %s", config_filename) + with open(config_filename, "r") as f: + text_format.Parse(f.read(), config) + + model = _load_model_from_checkpoint(config) + dl_config = _get_dataloader_config(config, batch_size_override) + evaluation = Evaluation(loss_fn=LczeroLoss(config=config.training.losses)) + dumper = Dumper(dump_to_stdout, dump_to_file, dump_to_shelve, dump_to_json) + onnx_comparator = OnnxComparator(onnx_model) if onnx_model else None + + samples_to_process = num_samples if num_samples is not None else 10 + logger.info("Starting evaluation with %d samples", samples_to_process) + + try: + evaluation.run( + model=model, + datagen=from_dataloader(make_dataloader(dl_config)), + num_samples=samples_to_process, + dumper=dumper, + onnx_comparator=onnx_comparator, + softmax_jax_wdl=softmax_jax_wdl, + ) + finally: + dumper.close() + if onnx_comparator: + onnx_comparator.log_summary() + + logger.info("Evaluation complete") diff --git a/src/lczero_training/training/init.py b/src/lczero_training/training/init.py new file mode 100644 index 00000000..dbf0a0a1 --- /dev/null +++ b/src/lczero_training/training/init.py @@ -0,0 +1,163 @@ +import gzip +import logging +import os +import sys +from typing import Optional + +import orbax.checkpoint as ocp +from flax import nnx +from google.protobuf import text_format + +from lczero_training.convert.leela_to_jax import ( + LeelaImportOptions, + fix_older_weights_file, + leela_to_jax, +) +from lczero_training.convert.leela_to_modelconfig import leela_to_modelconfig +from lczero_training.training.state import TrainingState +from proto import hlo_pb2, net_pb2 +from proto.model_config_pb2 import ModelConfig +from proto.root_config_pb2 import RootConfig + +logger = logging.getLogger(__name__) + + +def _load_lc0_model_state( + path: str, + expected_config: ModelConfig, + compute_dtype: hlo_pb2.XlaShapeProto.Type, + ignore_config_mismatch: bool = False, +) -> tuple[nnx.State, int]: + """Load lc0 weights, validate config, return (model_state, training_steps).""" + lc0_weights = net_pb2.Net() + with gzip.open(path, "rb") as f: + lc0_weights.ParseFromString(f.read()) + fix_older_weights_file(lc0_weights) + + leela_config = leela_to_modelconfig( + lc0_weights, hlo_pb2.XlaShapeProto.F32, compute_dtype + ) + if leela_config != expected_config: + if ignore_config_mismatch: + logger.warning( + "The provided lczero model configuration " + "differs from the one in the config file (ignored)." + ) + else: + logger.error( + "The provided lczero model configuration " + "differs from the one in the config file." + ) + logger.error(f"Config file model config: {expected_config}") + logger.error(f"Leela model config: {leela_config}") + sys.exit(1) + + import_options = LeelaImportOptions( + weights_dtype=hlo_pb2.XlaShapeProto.F32, compute_dtype=compute_dtype + ) + model_state = leela_to_jax(lc0_weights, import_options) + return model_state, lc0_weights.training_params.training_steps + + +def init( + config_filename: str, + lczero_model: Optional[str], + seed: int = 42, + dry_run: bool = False, + swa_initial_nets: int = 0, + override_training_steps: Optional[int] = None, + overwrite: bool = False, + no_copy_swa: bool = False, + ignore_config_mismatch: bool = False, +) -> None: + """ + Initializes a new training run. + """ + + config = RootConfig() + logger.info("Reading configuration from proto file") + with open(config_filename, "r") as f: + text_format.Parse(f.read(), config) + + checkpoint_path = config.training.checkpoint.path + checkpoint_exists = checkpoint_path and os.path.exists(checkpoint_path) + + if not dry_run and checkpoint_exists and not overwrite: + logger.error(f"Checkpoint path {checkpoint_path} already exists.") + sys.exit(1) + + logger.info("Creating initial training state from configuration") + training_state = TrainingState.new_from_config( + model_config=config.model, + training_config=config.training, + ) + + if checkpoint_exists: + logger.info(f"Loading from existing checkpoint: {checkpoint_path}") + source_mgr = ocp.CheckpointManager( + checkpoint_path, + options=ocp.CheckpointManagerOptions(create=False), + ) + training_state = source_mgr.restore( + source_mgr.latest_step(), + args=ocp.args.PyTreeRestore(training_state), + ) + + swa_enabled = config.training.HasField("swa") + + if lczero_model is None: + if override_training_steps is not None: + training_state = training_state.with_updated_step( + override_training_steps + ) + if swa_enabled and swa_initial_nets > 0: + training_state = training_state.replace( + jit_state=training_state.jit_state.replace( + num_averages=float(swa_initial_nets), + ) + ) + else: + logger.info(f"Loading lczero model: {lczero_model}") + model_state, lc0_steps = _load_lc0_model_state( + lczero_model, + config.model, + config.model.defaults.compute_dtype, + ignore_config_mismatch, + ) + step = override_training_steps or lc0_steps + new_swa_state = ( + training_state.jit_state.swa_state + if no_copy_swa + else (model_state if swa_enabled else None) + ) + training_state = training_state.replace( + jit_state=training_state.jit_state.replace( + model_state=model_state, + swa_state=new_swa_state, + num_averages=float(swa_initial_nets) if swa_enabled else 0.0, + ) + ).with_updated_step(step) + + if dry_run: + logger.info( + f"Would save checkpoint to {config.training.checkpoint.path} " + f"at step {training_state.jit_state.step}" + ) + else: + checkpoint_mgr = ocp.CheckpointManager( + config.training.checkpoint.path, + options=ocp.CheckpointManagerOptions(create=True), + ) + + step = training_state.jit_state.step + if step in checkpoint_mgr.all_steps(): + logger.info(f"Deleting existing checkpoint at step {step}") + checkpoint_mgr.delete(step) + checkpoint_mgr.wait_until_finished() + + logger.info( + f"Saving checkpoint to {config.training.checkpoint.path} at step {step}" + ) + checkpoint_mgr.save(step=step, args=ocp.args.PyTreeSave(training_state)) + checkpoint_mgr.wait_until_finished() + logger.info("Initialization complete.") diff --git a/src/lczero_training/training/lr_schedule.py b/src/lczero_training/training/lr_schedule.py new file mode 100644 index 00000000..a330ebd7 --- /dev/null +++ b/src/lczero_training/training/lr_schedule.py @@ -0,0 +1,145 @@ +from typing import Callable, Sequence + +import jax.numpy as jnp +import optax + +from proto.training_config_pb2 import LrSchedule + + +def _create_rule_fn(rule: LrSchedule) -> Callable: + """ + Creates a JAX-compatible function for a single LR schedule rule. + + All data from the protobuf is extracted here and captured by the closure of + the returned function. This avoids protobuf parsing inside the main schedule + function which will be JIT-compiled. + """ + start_step = float(rule.starting_step) + durations = list(rule.duration_steps) + lrs = list(rule.lr) + is_looping = rule.loop + + # Handle simple cases where the LR is constant for this rule. + if not durations or not lrs: + lr_val = jnp.asarray(lrs[-1] if lrs else 0.0, dtype=jnp.float32) + # Return a simple lambda that ignores the step and returns the constant value. + return lambda step: lr_val + + period = sum(durations) + if period == 0.0: + lr_val = jnp.asarray(lrs[-1], dtype=jnp.float32) + return lambda step: lr_val + + # Pre-calculate JAX arrays for use in the schedule function. + transitions = [ + ( + rule.transition[i] + if i < len(rule.transition) + else LrSchedule.Transition.CONSTANT + ) + for i in range(len(durations)) + ] + + durs_j = jnp.asarray(durations, dtype=jnp.float32) + ends_j = jnp.cumsum(durs_j) + starts_j = ends_j - durs_j + + # Interpolation start/end LRs for each segment. + a_vals = [ + lrs[i] if i < len(lrs) else lrs[-1] for i in range(len(durations)) + ] + b_vals = [ + lrs[i + 1] if (i + 1) < len(lrs) else a_vals[i] + for i in range(len(durations)) + ] + lrs_a_j = jnp.asarray(a_vals, dtype=jnp.float32) + lrs_b_j = jnp.asarray(b_vals, dtype=jnp.float32) + + trans_j = jnp.asarray(transitions, dtype=jnp.int32) + last_lr_j = jnp.asarray(lrs[-1], dtype=jnp.float32) + period_j = jnp.asarray(period, dtype=jnp.float32) + + def rule_fn(step: jnp.ndarray) -> jnp.ndarray: + """JAX-compatible function evaluating the LR for a given step.""" + rel_step = step - start_step + + if is_looping: + rel_step = jnp.mod(rel_step, period_j) + + # Find active segment and calculate interpolation factor `t`. + is_in_segment = ( + (rel_step >= starts_j) & (rel_step < ends_j) & (durs_j > 0) + ) + # Use maximum() to avoid division by zero for zero-duration segments. + t = jnp.clip((rel_step - starts_j) / jnp.maximum(durs_j, 1.0), 0.0, 1.0) + + # Calculate interpolated values for all segments for all transition types. + lin = lrs_a_j + (lrs_b_j - lrs_a_j) * t + cos = lrs_a_j + 0.5 * (1.0 - jnp.cos(jnp.pi * t)) * (lrs_b_j - lrs_a_j) + + # Select interpolation type for each segment. Default to CONSTANT (lrs_a_j). + interp_vals = jnp.where( + trans_j == LrSchedule.Transition.LINEAR, lin, lrs_a_j + ) + interp_vals = jnp.where( + trans_j == LrSchedule.Transition.COSINE, cos, interp_vals + ) + + # Select the value from the active segment by masking. + lr = jnp.sum(interp_vals * is_in_segment) + + # If not in any segment (e.g., gap between segments), use the last LR value. + lr = jnp.where(jnp.any(is_in_segment), lr, last_lr_j) + + if not is_looping: + # For non-looping rules, if past the end, clamp to the last LR. + lr = jnp.where(rel_step >= period_j, last_lr_j, lr) + + return lr + + return rule_fn + + +def make_lr_schedule(schedules: Sequence[LrSchedule]) -> optax.Schedule: + """ + Creates a learning rate schedule from a sequence of LrSchedule protobufs. + + The schedule is composed of multiple rules, each active for a certain + range of training steps. + """ + if not schedules: + return lambda count: jnp.asarray(0.0, dtype=jnp.float32) + + rule_fns = [_create_rule_fn(rule) for rule in schedules] + start_steps = jnp.asarray( + [rule.starting_step for rule in schedules], dtype=jnp.float32 + ) + + # Determine the learning rate to use for steps before the first rule begins. + first_lrs = jnp.asarray( + [rule.lr[0] if rule.lr else 0.0 for rule in schedules], + dtype=jnp.float32, + ) + earliest_rule_idx = jnp.argmin(start_steps) + pre_start_lr = first_lrs[earliest_rule_idx] + min_start_step = start_steps[earliest_rule_idx] + + def schedule(count: jnp.ndarray) -> jnp.ndarray: + """The actual schedule function passed to Optax.""" + step = jnp.asarray(count, dtype=jnp.float32) + + # Find the index of the active rule. The active rule is the one with the + # largest starting_step that is less than or equal to the current step. + eligible_mask = step >= start_steps + # Replace non-eligible start_steps with a large negative number so they + # are ignored by argmax. + effective_starts = jnp.where(eligible_mask, start_steps, -1.0) + active_rule_idx = jnp.argmax(effective_starts) + + all_lrs = jnp.stack([fn(step) for fn in rule_fns]) + lr = all_lrs[active_rule_idx] + + # If the current step is before any rule starts, use the pre-start LR. + return jnp.where(step < min_start_step, pre_start_lr, lr) + + return schedule diff --git a/src/lczero_training/training/migrate_checkpoint.py b/src/lczero_training/training/migrate_checkpoint.py new file mode 100644 index 00000000..824266cc --- /dev/null +++ b/src/lczero_training/training/migrate_checkpoint.py @@ -0,0 +1,345 @@ +from typing import Any, Dict, Iterable, List, Set, Tuple + +import jax +import numpy as np +import orbax.checkpoint as ocp +from flax import serialization +from google.protobuf import text_format +from orbax.checkpoint.utils import tuple_path_from_keypath + +from lczero_training.training import state as state_lib +from proto import checkpoint_migration_config_pb2, root_config_pb2 + + +def _str_to_key_path(path_str: str) -> tuple[str, ...]: + return tuple(path_str.split(".")) + + +def _load_new_state( + root_config: root_config_pb2.RootConfig, serialized_model: bool +) -> Any: + new_state = state_lib.TrainingState.new_from_config( + root_config.model, root_config.training + ) + if serialized_model: + return serialization.to_state_dict(new_state) + return new_state + + +def load_checkpoint( + checkpoint_path: str, checkpoint_step: int | None = None +) -> Tuple[Any, int]: + """Load a checkpoint from the given path. + + Args: + checkpoint_path: Path to the checkpoint directory. + checkpoint_step: Step to load, or None to load the latest. + + Returns: + Tuple of (checkpoint_state, checkpoint_step). + """ + manager = ocp.CheckpointManager( + checkpoint_path, + options=ocp.CheckpointManagerOptions(create=False), + ) + if checkpoint_step is None: + checkpoint_step = manager.latest_step() + if checkpoint_step is None: + raise ValueError(f"No checkpoints found in {checkpoint_path}") + return manager.restore(checkpoint_step), checkpoint_step + + +def get_checkpoint_steps( + checkpoint_path: str, + min_step: int | None = None, + max_step: int | None = None, +) -> list[int]: + """Get all checkpoint steps in the given range. + + Args: + checkpoint_path: Path to the checkpoint directory. + min_step: Minimum step (inclusive), or None for no minimum. + max_step: Maximum step (inclusive), or None for no maximum. + + Returns: + List of checkpoint steps in ascending order. + """ + manager = ocp.CheckpointManager( + checkpoint_path, + options=ocp.CheckpointManagerOptions(create=False), + ) + all_steps = sorted(manager.all_steps()) + filtered_steps = [] + for step in all_steps: + if min_step is not None and step < min_step: + continue + if max_step is not None and step > max_step: + continue + filtered_steps.append(step) + filtered_steps.sort() + return filtered_steps + + +def _load_old_state( + checkpoint_path: str, checkpoint_step: int | None +) -> Tuple[Any, int]: + return load_checkpoint(checkpoint_path, checkpoint_step) + + +def load_migration_rules(rules_file: str | None) -> List[Tuple[Any, Any]]: + """Load migration rules from a CheckpointMigrationConfig file. + + Args: + rules_file: Path to the CheckpointMigrationConfig textproto file, + or None to return empty rules. + + Returns: + List of (from_path, to_path) tuples representing migration rules. + """ + rules = [] + if rules_file: + migration_config = ( + checkpoint_migration_config_pb2.CheckpointMigrationConfig() + ) + with open(rules_file, "r") as f: + text_format.Parse(f.read(), migration_config) + for rule_proto in migration_config.rule: + from_path = ( + _str_to_key_path(rule_proto.from_path) + if rule_proto.from_path + else None + ) + to_path = ( + _str_to_key_path(rule_proto.to_path) + if rule_proto.to_path + else None + ) + rules.append((from_path, to_path)) + return rules + + +def _format_value(value: Any) -> str: + if isinstance(value, (np.ndarray, jax.Array)): + return f"{value.dtype}{value.shape}" + return repr(value) + + +def _format_path_diff( + unhandled_source: Set[Tuple[str, ...]], + unhandled_dest: Set[Tuple[str, ...]], + old_paths: Dict[Tuple[str, ...], Any], + new_paths: Dict[Tuple[str, ...], Any], +) -> str: + diff = [] + for p in sorted(list(unhandled_source | unhandled_dest)): + p_str = ".".join(p) + if p in unhandled_source: + diff.append(f"- {p_str}: {_format_value(old_paths[p])}") + if p in unhandled_dest: + diff.append(f"+ {p_str}: {_format_value(new_paths[p])}") + return "\n".join(diff) + + +class Migration: + def __init__(self, old_state: Any, new_state: Any): + old_leaves = jax.tree_util.tree_leaves_with_path(old_state) + new_flat, self.new_treedef = jax.tree_util.tree_flatten_with_path( + new_state + ) + + self.old_paths: Dict[Tuple[str, ...], Any] = { + tuple_path_from_keypath(path): value for path, value in old_leaves + } + self.new_paths: Dict[Tuple[str, ...], Any] = { + tuple_path_from_keypath(path): value for path, value in new_flat + } + self.new_leaves: List[Any] = [value for _, value in new_flat] + self.new_path_to_idx: Dict[Tuple[str, ...], int] = { + tuple_path_from_keypath(path): i + for i, (path, _) in enumerate(new_flat) + } + + self.source_paths: Set[Tuple[str, ...]] = set(self.old_paths.keys()) + self.dest_paths: Set[Tuple[str, ...]] = set(self.new_path_to_idx.keys()) + + print(f"{len(self.source_paths & self.dest_paths)} common keys") + print(f"{len(self.source_paths - self.dest_paths)} keys disappeared") + print(f"{len(self.dest_paths - self.source_paths)} keys appeared") + + self.errors: List[str] = [] + + def _apply_move_rule( + self, from_path: Tuple[str, ...], to_path: Tuple[str, ...] + ) -> None: + if from_path == to_path: + self.errors.append( + f"from_path and to_path are the same: {from_path}" + ) + return + + source_prefixed = { + p for p in self.source_paths if p[: len(from_path)] == from_path + } + dest_prefixed = { + p for p in self.dest_paths if p[: len(to_path)] == to_path + } + + if not source_prefixed: + self.errors.append(f"from_path {from_path} not found in old state") + if not dest_prefixed: + self.errors.append(f"to_path {to_path} not found in new state") + + for p in source_prefixed: + new_p = to_path + p[len(from_path) :] + if new_p in self.dest_paths: + idx = self.new_path_to_idx[new_p] + self.new_leaves[idx] = self.old_paths[p] + self.source_paths.remove(p) + self.dest_paths.remove(new_p) + else: + self.errors.append(f"Path {new_p} not found in new state") + + def _apply_ignore_rule(self, from_path: Tuple[str, ...]) -> None: + source_prefixed = { + p for p in self.source_paths if p[: len(from_path)] == from_path + } + if not source_prefixed: + self.errors.append(f"from_path {from_path} not found in old state") + self.source_paths -= source_prefixed + + def _apply_keep_rule(self, to_path: Tuple[str, ...]) -> None: + dest_prefixed = { + p for p in self.dest_paths if p[: len(to_path)] == to_path + } + if not dest_prefixed: + self.errors.append(f"to_path {to_path} not found in new state") + self.dest_paths -= dest_prefixed + + def apply_rules(self, rules: List[Tuple[Any, Any]]) -> None: + for from_path, to_path in rules: + if from_path and to_path: + self._apply_move_rule(from_path, to_path) + elif from_path: + self._apply_ignore_rule(from_path) + elif to_path: + self._apply_keep_rule(to_path) + + def run(self, rules: List[Tuple[Any, Any]]) -> Any: + self.apply_rules(rules) + + # Copy remaining paths + copied_paths = self.source_paths & self.dest_paths + for p in copied_paths: + idx = self.new_path_to_idx[p] + self.new_leaves[idx] = self.old_paths[p] + self.source_paths -= copied_paths + self.dest_paths -= copied_paths + + unhandled_source = self.source_paths + unhandled_dest = self.dest_paths + + if unhandled_source or unhandled_dest: + self.errors.append( + "Unmapped paths:\n" + + _format_path_diff( + unhandled_source, + unhandled_dest, + self.old_paths, + self.new_paths, + ) + ) + + if self.errors: + raise ValueError("\n".join(self.errors)) + + return getattr(self.new_treedef, "unflatten")(self.new_leaves) + + +def _save_checkpoint( + migrated_state: Any, + new_checkpoint_path: str, + new_checkpoint_step: int, + overwrite: bool, +) -> None: + manager = ocp.CheckpointManager( + new_checkpoint_path, + ocp.PyTreeCheckpointer(), + options=ocp.CheckpointManagerOptions( + create=True, save_interval_steps=1, todelete_subdir="trash" + ), + ) + + if new_checkpoint_step in manager.all_steps(): + if overwrite: + manager.delete(new_checkpoint_step) + manager.wait_until_finished() + else: + raise ValueError( + f"Checkpoint already exists at {new_checkpoint_step} in " + f"{new_checkpoint_path}. " + "Use --overwrite to overwrite." + ) + + manager.save(new_checkpoint_step, migrated_state) + manager.wait_until_finished() + print( + f"New checkpoint saved successfully to {new_checkpoint_path} at step " + f"{new_checkpoint_step}." + ) + + +def _dump_paths(paths: Iterable[Tuple[str, ...]], field: str) -> None: + def key(p: Tuple[str, ...]) -> tuple: + return tuple(int(c) if c.isdigit() else c for c in p) + + for path in sorted(paths, key=key): + print(f'rule {{ {field}: "{".".join(path)}" }}') + + +def migrate_checkpoint( + config: str, + new_checkpoint: str | None, + overwrite: bool, + rules_file: str | None, + serialized_model: bool, + checkpoint_step: int | None, + new_checkpoint_step: int | None, + dump_source_paths: bool = False, + dump_destination_paths: bool = False, +) -> None: + """Migrates a checkpoint to a new training state.""" + root_config = root_config_pb2.RootConfig() + with open(config, "r") as f: + text_format.Parse(f.read(), root_config) + + new_state = _load_new_state(root_config, serialized_model) + old_state, old_checkpoint_step = _load_old_state( + root_config.training.checkpoint.path, checkpoint_step + ) + rules = load_migration_rules(rules_file) + + migration = Migration(old_state, new_state) + + if dump_source_paths: + _dump_paths(migration.old_paths.keys(), "from_path") + if dump_destination_paths: + _dump_paths(migration.new_paths.keys(), "to_path") + if dump_source_paths or dump_destination_paths: + return + + migrated_state = migration.run(rules) + + if new_checkpoint: + checkpoint_path = new_checkpoint + elif overwrite: + checkpoint_path = root_config.training.checkpoint.path + else: + print("Migration check successful.") + return + + if new_checkpoint_step is None: + new_checkpoint_step = old_checkpoint_step + + _save_checkpoint( + migrated_state, checkpoint_path, new_checkpoint_step, overwrite + ) diff --git a/src/lczero_training/training/optimizer.py b/src/lczero_training/training/optimizer.py new file mode 100644 index 00000000..16eaed3c --- /dev/null +++ b/src/lczero_training/training/optimizer.py @@ -0,0 +1,86 @@ +from functools import partial + +import jax +import jax.numpy as jnp +import optax +from flax import nnx + +from lczero_training.training.utils import make_weights_mask +from proto.training_config_pb2 import OptimizerConfig + +_STATES_WITH_COUNT = ( + optax.ScaleByAdamState, + optax.ScaleByScheduleState, +) + + +def update_optimizer_step( + opt_state: optax.OptState, step: int +) -> optax.OptState: + """Updates all step counters in the optimizer state tree.""" + step_array = jnp.array(step, dtype=jnp.int32) + + def has_count(x: object) -> bool: + return hasattr(x, "_fields") and "count" in x._fields + + def update_count(x: optax.OptState) -> optax.OptState: + if not has_count(x): + return x + assert isinstance(x, _STATES_WITH_COUNT), ( + f"Unexpected state type with 'count' field: {type(x).__name__}" + ) + return x._replace(count=step_array) + + return jax.tree_util.tree_map(update_count, opt_state, is_leaf=has_count) + + +def make_gradient_transformation( + config: OptimizerConfig, + *, + max_grad_norm: float | None = None, + lr_schedule: optax.Schedule, +) -> optax.GradientTransformation: + if config.HasField("nadamw"): + nadamw = config.nadamw + tx = optax.nadamw( + lr_schedule, + b1=nadamw.beta_1, + b2=nadamw.beta_2, + eps=nadamw.epsilon, + weight_decay=nadamw.weight_decay, + mask=partial(make_weights_mask, nadamw.decay_selector), + ) + elif config.HasField("nadam"): + nadam = config.nadam + tx = optax.nadam( + lr_schedule, + b1=nadam.beta_1, + b2=nadam.beta_2, + eps=nadam.epsilon, + ) + elif config.HasField("sgd"): + sgd = config.sgd + tx = optax.sgd( + lr_schedule, + momentum=sgd.momentum if sgd.momentum else None, + nesterov=sgd.nesterov, + ) + else: + raise ValueError( + "Unsupported optimizer type: {}".format( + config.WhichOneof("optimizer_type") + ) + ) + if max_grad_norm is not None and max_grad_norm > 0: + tx = optax.chain(optax.clip_by_global_norm(max_grad_norm), tx) + if config.HasField("freeze_selector"): + freeze_mask = partial(make_weights_mask, config.freeze_selector) + + def trainable_mask(p: nnx.State) -> nnx.State: + return jax.tree.map(lambda x: not x, freeze_mask(p)) + + tx = optax.chain( + optax.masked(tx, trainable_mask), + optax.masked(optax.set_to_zero(), freeze_mask), + ) + return tx diff --git a/src/lczero_training/training/overfit.py b/src/lczero_training/training/overfit.py new file mode 100644 index 00000000..134549fb --- /dev/null +++ b/src/lczero_training/training/overfit.py @@ -0,0 +1,289 @@ +"""Overfitting utility for quickly validating training setup.""" + +import csv +import logging +from contextlib import suppress +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx +from google.protobuf import text_format +from jax import tree_util + +from lczero_training.dataloader import DataLoader, make_dataloader +from lczero_training.model.loss_function import LczeroLoss +from lczero_training.model.model import LczeroModel +from lczero_training.training.lr_schedule import make_lr_schedule +from lczero_training.training.optimizer import make_gradient_transformation +from lczero_training.training.state import ( + TrainingBatch, + TrainingSample, + TrainingState, +) +from lczero_training.training.training import Training +from proto.root_config_pb2 import RootConfig + +logger = logging.getLogger(__name__) + + +def _stop_loader(loader: DataLoader) -> None: + with suppress(Exception): + loader.stop() + + +def _prepare_batch(batch_tuple: tuple) -> TrainingBatch: + # DataLoader now returns tuple: (inputs, probabilities, values) + return TrainingBatch( + inputs=jnp.asarray(batch_tuple[0]), + probabilities=jnp.asarray(batch_tuple[1]), + values=jnp.asarray(batch_tuple[2]), + ) + + +def _make_eval_step(graphdef: nnx.GraphDef, loss_fn: LczeroLoss) -> Any: + @partial(nnx.jit, static_argnames=()) + def eval_step( + model_state: nnx.State, batch: TrainingBatch + ) -> tuple[jax.Array, Any]: + model = nnx.merge(graphdef, model_state) + + def loss_for_batch( + model_arg: LczeroModel, sample_arg: TrainingSample + ) -> tuple[jax.Array, Any]: + return loss_fn(model_arg, sample_arg) + + loss_vfn = jax.vmap(loss_for_batch, in_axes=(None, 0), out_axes=0) + # vmap automatically distributes TrainingBatch over batch dimension, + # calling loss_for_batch with TrainingSample (single samples). + per_sample_loss, unweighted_losses = loss_vfn(model, batch) # type: ignore[arg-type] + mean_loss = jnp.mean(per_sample_loss) + mean_unweighted = tree_util.tree_map(jnp.mean, unweighted_losses) + return mean_loss, mean_unweighted + + return eval_step + + +def overfit( + *, + config_filename: str, + num_steps: int, + coin_flip: bool = False, + csv_file: str | None = None, +) -> None: + """Runs an overfitting loop to validate training.""" + + if num_steps <= 0: + raise ValueError("num_steps must be a positive integer") + + if jax.device_count() > 1: + raise ValueError( + f"Overfit utility does not support multi-GPU training. " + f"Detected {jax.device_count()} devices. " + f"Please set CUDA_VISIBLE_DEVICES to use only one GPU." + ) + + config = RootConfig() + logger.info("Reading configuration from proto file") + with open(config_filename, "r") as config_file: + text_format.Parse(config_file.read(), config) + + logger.info("Creating data loader and fetching batches") + loader = make_dataloader(config.data_loader) + try: + batch_a = loader.get_next() + batch_b = loader.get_next() if coin_flip else None + finally: + _stop_loader(loader) + + prepared_batch_a = _prepare_batch(batch_a) + prepared_batch_b = _prepare_batch(batch_b) if batch_b is not None else None + + logger.info("Creating training state from configuration") + training_state = TrainingState.new_from_config( + model_config=config.model, + training_config=config.training, + ) + + graphdef, _ = nnx.split( + LczeroModel(config=config.model, rngs=nnx.Rngs(params=42)) + ) + + jit_state = training_state.jit_state + lr_sched = make_lr_schedule(config.training.lr_schedule) + optimizer_tx = make_gradient_transformation( + config.training.optimizer, + max_grad_norm=getattr(config.training, "max_grad_norm", 0.0), + lr_schedule=lr_sched, + ) + + loss_fn = LczeroLoss(config=config.training.losses) + training = Training( + optimizer_tx=optimizer_tx, + graphdef=graphdef, + loss_fn=loss_fn, + ) + eval_step = _make_eval_step(graphdef, loss_fn) + + csv_handle = None + csv_writer: Any | None = None + if csv_file is not None: + logger.info("Writing overfit results to %s", csv_file) + csv_handle = open(csv_file, "w", newline="") + csv_writer = csv.writer(csv_handle) + csv_writer.writerow( + [ + "step", + "train_batch", + "train_loss", + "train_unweighted", + "eval_batch", + "eval_loss", + "eval_unweighted", + ] + ) + csv_handle.flush() + + def log_step( + *, + step_value: int, + train_batch_name: str, + train_loss: float, + train_unweighted: Any, + eval_batch_name: str | None, + eval_loss: float | None, + eval_unweighted: Any | None, + ) -> None: + if eval_batch_name is None or eval_loss is None: + logger.info( + "Step %d: batch=%s train_loss=%f, unweighted_losses=%s", + step_value, + train_batch_name, + train_loss, + train_unweighted, + ) + else: + logger.info( + ( + "Step %d: trained %s train_loss=%f, unweighted_losses=%s; " + "evaluated %s eval_loss=%f, eval_unweighted=%s" + ), + step_value, + train_batch_name, + train_loss, + train_unweighted, + eval_batch_name, + eval_loss, + eval_unweighted, + ) + + if csv_writer is not None and csv_handle is not None: + csv_writer.writerow( + [ + step_value, + train_batch_name, + train_loss, + repr(train_unweighted), + eval_batch_name or "", + "" if eval_loss is None else eval_loss, + "" if eval_unweighted is None else repr(eval_unweighted), + ] + ) + csv_handle.flush() + + try: + if coin_flip: + if prepared_batch_b is None: + raise RuntimeError( + "Coin flip mode requires two batches but only one was fetched" + ) + + logger.info( + "Starting coin-flip overfit: %d steps on batch A then %d on batch B", + num_steps, + num_steps, + ) + + def run_phase( + train_batch: TrainingBatch, + train_name: str, + eval_batch: TrainingBatch, + eval_name: str, + ) -> None: + nonlocal jit_state + for _ in range(num_steps): + jit_state, metrics = training.train_step( + optimizer_tx, + jit_state, + train_batch, + ) + loss = metrics["loss"] + unweighted_losses = metrics["unweighted_losses"] + loss_value, unweighted_host = jax.device_get( + (loss, unweighted_losses) + ) + loss_value = float(np.asarray(loss_value)) + unweighted_host = tree_util.tree_map( + lambda x: float(np.asarray(x)), unweighted_host + ) + + eval_loss, eval_unweighted = eval_step( + jit_state.model_state, eval_batch + ) + eval_loss, eval_unweighted = jax.device_get( + (eval_loss, eval_unweighted) + ) + eval_loss_value = float(np.asarray(eval_loss)) + eval_unweighted_host = tree_util.tree_map( + lambda x: float(np.asarray(x)), eval_unweighted + ) + + step_value = int( + np.asarray(jax.device_get(jit_state.step)).flat[0] + ) + log_step( + step_value=step_value, + train_batch_name=train_name, + train_loss=loss_value, + train_unweighted=unweighted_host, + eval_batch_name=eval_name, + eval_loss=eval_loss_value, + eval_unweighted=eval_unweighted_host, + ) + + run_phase(prepared_batch_a, "A", prepared_batch_b, "B") + run_phase(prepared_batch_b, "B", prepared_batch_a, "A") + else: + logger.info("Starting overfit loop for %d steps", num_steps) + for _ in range(num_steps): + jit_state, metrics = training.train_step( + optimizer_tx, + jit_state, + prepared_batch_a, + ) + loss = metrics["loss"] + unweighted_losses = metrics["unweighted_losses"] + loss_value, unweighted_host = jax.device_get( + (loss, unweighted_losses) + ) + loss_value = float(np.asarray(loss_value)) + unweighted_host = tree_util.tree_map( + lambda x: float(np.asarray(x)), unweighted_host + ) + step_value = int( + np.asarray(jax.device_get(jit_state.step)).flat[0] + ) + log_step( + step_value=step_value, + train_batch_name="single", + train_loss=loss_value, + train_unweighted=unweighted_host, + eval_batch_name=None, + eval_loss=None, + eval_unweighted=None, + ) + finally: + if csv_handle is not None: + csv_handle.close() diff --git a/src/lczero_training/training/state.py b/src/lczero_training/training/state.py new file mode 100644 index 00000000..62f12f42 --- /dev/null +++ b/src/lczero_training/training/state.py @@ -0,0 +1,149 @@ +import dataclasses +import logging +from typing import Any, Optional, Union + +import jax +import jax.numpy as jnp +import jax.sharding as jshard +import numpy as np +import optax +from flax import nnx +from flax.struct import dataclass + +from lczero_training.model.model import LczeroModel +from lczero_training.training.lr_schedule import make_lr_schedule +from lczero_training.training.optimizer import ( + make_gradient_transformation, + update_optimizer_step, +) +from proto.model_config_pb2 import ModelConfig +from proto.training_config_pb2 import TrainingConfig + +logger = logging.getLogger(__name__) + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass +class TrainingSample: + """Single training sample without batch dimension. + + Used for vmap over individual samples in loss computation. + + Fields: + inputs: Input planes tensor [112, 8, 8] + probabilities: Policy probabilities tensor [1858] + values: Combined values tensor [6, 3] where: + - Index 0: result [result_q, result_d, plies_left] + - Index 1: best [best_q, best_d, best_m] + - Index 2: played [played_q, played_d, played_m] + - Index 3: orig [orig_q, orig_d, orig_m] (may contain NaN) + - Index 4: root [root_q, root_d, root_m] + - Index 5: st [q_st, d_st, NaN] + """ + + inputs: jax.Array + probabilities: jax.Array + values: jax.Array + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass +class TrainingBatch: + """Batch of training data with inputs, probabilities, and values tensors. + + Fields: + inputs: Input planes tensor [batch, 112, 8, 8] + probabilities: Policy probabilities tensor [batch, 1858] + values: Combined values tensor [batch, 6, 3] where: + - Index 0: result [result_q, result_d, plies_left] + - Index 1: best [best_q, best_d, best_m] + - Index 2: played [played_q, played_d, played_m] + - Index 3: orig [orig_q, orig_d, orig_m] (may contain NaN) + - Index 4: root [root_q, root_d, root_m] + - Index 5: st [q_st, d_st, NaN] + """ + + inputs: Union[jax.Array, jshard.NamedSharding] + probabilities: Union[jax.Array, jshard.NamedSharding] + values: Union[jax.Array, jshard.NamedSharding] + + @classmethod + def from_tuple( + cls, tensor_tuple: tuple[np.ndarray, ...] + ) -> "TrainingBatch": + """Create TrainingBatch from tuple returned by DataLoader.""" + if len(tensor_tuple) != 3: + raise ValueError( + f"Expected tuple of 3 tensors, got {len(tensor_tuple)}" + ) + return cls( + inputs=jnp.asarray(tensor_tuple[0]), + probabilities=jnp.asarray(tensor_tuple[1]), + values=jnp.asarray(tensor_tuple[2]), + ) + + +@dataclass +class JitTrainingState: + step: int + model_state: nnx.State + opt_state: Optional[optax.OptState] + # SWA state mirrors model_state structure when enabled; None otherwise. + # Marked non-pytree to exclude from JIT/pjit inputs and device transfers. + swa_state: Optional[nnx.State] + # Effective number of model snapshots accumulated into SWA (can be fractional). + num_averages: float + + def replace(self, **changes: Any) -> "JitTrainingState": + """Returns a new instance of the class with the specified changes.""" + return dataclasses.replace(self, **changes) + + +@dataclass +class TrainingState: + jit_state: JitTrainingState + # Last chunk source that was available when the last epoch started training. + num_heads: int + last_chunk_source: str = "" + + def replace(self, **changes: Any) -> "TrainingState": + """Returns a new instance of the class with the specified changes.""" + return dataclasses.replace(self, **changes) + + def with_updated_step(self, step: int) -> "TrainingState": + """Returns a copy with updated step in both jit_state and optimizer.""" + updated_opt_state = ( + update_optimizer_step(self.jit_state.opt_state, step) + if self.jit_state.opt_state is not None + else None + ) + return self.replace( + jit_state=self.jit_state.replace( + step=step, + opt_state=updated_opt_state, + ) + ) + + @staticmethod + def new_from_config( + model_config: ModelConfig, training_config: TrainingConfig + ) -> "TrainingState": + rngs = nnx.Rngs(params=42) + model_state = nnx.state(LczeroModel(config=model_config, rngs=rngs)) + lr_sched = make_lr_schedule(training_config.lr_schedule) + opt_state = make_gradient_transformation( + training_config.optimizer, + max_grad_norm=getattr(training_config, "max_grad_norm", 0.0), + lr_schedule=lr_sched, + ).init(model_state) + jit_state = JitTrainingState( + step=0, + model_state=model_state, + opt_state=opt_state, + swa_state=model_state, + num_averages=0.0, + ) + return TrainingState( + jit_state=jit_state, + num_heads=model_config.encoder.heads, + ) diff --git a/src/lczero_training/training/tensorboard.py b/src/lczero_training/training/tensorboard.py new file mode 100644 index 00000000..97255c21 --- /dev/null +++ b/src/lczero_training/training/tensorboard.py @@ -0,0 +1,69 @@ +"""Utilities for writing training metrics to TensorBoard event files.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from typing import Any, Dict + +import jax +import numpy as np +from tensorboardX import SummaryWriter + +logger = logging.getLogger(__name__) + +MetricsDict = Dict[str, Any] + + +def _to_ndarray(value: Any) -> np.ndarray: + try: + return np.asarray(jax.device_get(value)) + except TypeError: + return np.asarray(value) + + +def _to_scalar(value: Any) -> float | None: + array = _to_ndarray(value) + if array.ndim == 0 or array.size == 1: + return float(array.reshape(())) + logger.warning( + "Skipping non-scalar metric with shape %s when logging to TensorBoard.", + array.shape, + ) + return None + + +def _flatten_metrics( + metrics: Mapping[str, Any], prefix: str = "" +) -> Dict[str, float]: + scalars: Dict[str, float] = {} + for key, value in metrics.items(): + tag = f"{prefix}{key}" if prefix else key + if isinstance(value, Mapping): + scalars.update(_flatten_metrics(value, f"{tag}/")) + continue + scalar = _to_scalar(value) + if scalar is not None: + scalars[tag] = scalar + return scalars + + +def _to_step(step: Any) -> int: + return int(_to_ndarray(step).reshape(())) + + +class TensorboardLogger: + """Writes scalar training metrics to TensorBoard.""" + + def __init__(self, logdir: str) -> None: + self._writer = SummaryWriter(logdir) + + def log(self, step: int, metrics: MetricsDict) -> None: + global_step = _to_step(step) + for tag, value in _flatten_metrics(metrics).items(): + self._writer.add_scalar( + tag=tag, scalar_value=value, global_step=global_step + ) + + def close(self) -> None: + self._writer.close() diff --git a/src/lczero_training/training/test_lr_schedule.py b/src/lczero_training/training/test_lr_schedule.py new file mode 100644 index 00000000..49692d05 --- /dev/null +++ b/src/lczero_training/training/test_lr_schedule.py @@ -0,0 +1,138 @@ +from typing import Callable, List + +import jax.numpy as jnp +import pytest + +from lczero_training.training.lr_schedule import make_lr_schedule +from proto import training_config_pb2 as pb + + +def _sched( + schedules: List[pb.LrSchedule], +) -> Callable[[jnp.ndarray], jnp.ndarray]: + return make_lr_schedule(schedules) + + +def _val(s: Callable[[jnp.ndarray], jnp.ndarray], t: int | float) -> float: + return float(s(jnp.asarray(t, dtype=jnp.float32))) + + +def test_rule_selection_by_starting_step() -> None: + r0 = pb.LrSchedule( + starting_step=0, + duration_steps=[5], + lr=[0.1, 0.2], + ) + r1 = pb.LrSchedule( + starting_step=10, + duration_steps=[5], + lr=[0.5, 0.5], + ) + sched = _sched([r0, r1]) + assert _val(sched, 0) == pytest.approx(0.1) + assert _val(sched, 9) == pytest.approx(0.2) + assert _val(sched, 10) == pytest.approx(0.5) + assert _val(sched, 100) == pytest.approx(0.5) + + +def test_default_constant_transition_and_tail() -> None: + r = pb.LrSchedule( + starting_step=0, + duration_steps=[5], + lr=[0.3, 0.8], # no transition specified -> CONSTANT + ) + sched = _sched([r]) + for t in range(5): + assert _val(sched, t) == pytest.approx(0.3) + # Beyond period (no loop) yields last lr + assert _val(sched, 6) == pytest.approx(0.8) + + +def test_linear_then_hold() -> None: + r = pb.LrSchedule( + starting_step=0, + duration_steps=[3, 7], + lr=[0.0, 0.9, 0.9], + transition=[pb.LrSchedule.Transition.LINEAR], + ) + sched = _sched([r]) + assert _val(sched, 0) == pytest.approx(0.0) + assert _val(sched, 1) == pytest.approx(0.3) + assert _val(sched, 2) == pytest.approx(0.6) + assert _val(sched, 3) == pytest.approx(0.9) + assert _val(sched, 8) == pytest.approx(0.9) + + +def test_looping_constant_segments() -> None: + r = pb.LrSchedule( + starting_step=0, + duration_steps=[3, 2], + lr=[1.0, 2.0, 3.0], + loop=True, + ) + sched = _sched([r]) + assert _val(sched, 0) == pytest.approx(1.0) + assert _val(sched, 2) == pytest.approx(1.0) + assert _val(sched, 3) == pytest.approx(2.0) + assert _val(sched, 5) == pytest.approx(1.0) # 5 % (3+2) == 0 + + +def test_zero_duration_is_skipped() -> None: + r = pb.LrSchedule( + starting_step=0, + duration_steps=[0, 5], + lr=[1.0, 2.0, 3.0], + transition=[ + pb.LrSchedule.Transition.LINEAR, + pb.LrSchedule.Transition.LINEAR, + ], + ) + sched = _sched([r]) + # Interpolates over the second interval [2.0 -> 3.0] + assert _val(sched, 0) == pytest.approx(2.0) + assert _val(sched, 2) == pytest.approx(2.4) + assert _val(sched, 5) == pytest.approx(3.0) + + +def test_chain_zero_durations_then_linear() -> None: + r = pb.LrSchedule( + starting_step=0, + duration_steps=[0, 0, 0, 5], + lr=[1.0, 2.0, 3.0, 4.0, 5.0], + transition=[ + pb.LrSchedule.Transition.LINEAR, + pb.LrSchedule.Transition.LINEAR, + pb.LrSchedule.Transition.LINEAR, + pb.LrSchedule.Transition.LINEAR, + ], + ) + sched = _sched([r]) + # Should use last interval [4.0 -> 5.0] linearly across 5 steps + assert _val(sched, 0) == pytest.approx(4.0) + assert _val(sched, 2) == pytest.approx(4.4) + assert _val(sched, 4) == pytest.approx(4.8) + + +def test_cosine() -> None: + r = pb.LrSchedule( + starting_step=0, + duration_steps=[4], + lr=[0.0, 1.0], + transition=[pb.LrSchedule.Transition.COSINE], + ) + sched = _sched([r]) + # t=0 -> 0.0, t=0.5 -> 0.5, t=1.0 -> 1.0 approximately + assert _val(sched, 0) == pytest.approx(0.0) + assert _val(sched, 2) == pytest.approx(0.5, abs=1e-6) + assert _val(sched, 4) == pytest.approx(1.0) + + +def test_before_first_rule_uses_earliest_first_lr() -> None: + r = pb.LrSchedule( + starting_step=5, + duration_steps=[3], + lr=[0.1, 0.2], + ) + sched = _sched([r]) + # Before the first rule starts, schedule returns the first lr of earliest rule. + assert _val(sched, 0) == pytest.approx(0.1) diff --git a/src/lczero_training/training/training.py b/src/lczero_training/training/training.py new file mode 100644 index 00000000..7f8e61bb --- /dev/null +++ b/src/lczero_training/training/training.py @@ -0,0 +1,329 @@ +import dataclasses +import logging +from datetime import datetime +from functools import partial +from typing import Any, Callable, Dict, Generator, Optional, Tuple, cast + +import jax +import jax.numpy as jnp +import jax.sharding as jshard +import numpy as np +import optax +from flax import nnx +from jax import tree_util +from jax.sharding import PartitionSpec as P + +from lczero_training.dataloader import DataLoader +from lczero_training.model.loss_function import LczeroLoss +from lczero_training.model.model import LczeroModel +from lczero_training.training.state import ( + JitTrainingState, + TrainingBatch, + TrainingSample, +) +from proto import training_config_pb2 as training_config_pb2 + +MetricsDict = Dict[str, Any] + + +@dataclasses.dataclass +class StepHookData: + """Data passed to the step hook callback during training.""" + + global_step: int + local_step: int + steps_per_epoch: int + metrics: MetricsDict + jit_state: JitTrainingState + + +StepHook = Callable[[StepHookData], None] + +logger = logging.getLogger(__name__) + + +def from_dataloader( + loader: DataLoader, +) -> Generator[tuple[np.ndarray, ...], None, None]: + while True: + yield loader.get_next() + + +class Training: + optimizer_tx: optax.GradientTransformation + train_step: Callable[ + [optax.GradientTransformation, JitTrainingState, TrainingBatch], + Tuple[JitTrainingState, MetricsDict], + ] + _swa_config: Optional[training_config_pb2.SWAConfig] + _dp_sharding: Optional[jshard.NamedSharding] + + def __init__( + self, + optimizer_tx: optax.GradientTransformation, + graphdef: nnx.GraphDef, + loss_fn: LczeroLoss, + swa_config: Optional[training_config_pb2.SWAConfig] = None, + ): + self.optimizer_tx = optimizer_tx + self._swa_config = swa_config + self._dp_sharding = None + + jit_kwargs: Dict[str, Any] = { + "static_argnames": ("optimizer_tx",), + "donate_argnames": ("jit_state",), + } + if jax.device_count() > 1: + num_devices = jax.device_count() + logger.info( + f"Multi-GPU training enabled: {num_devices} devices detected" + ) + mesh = jshard.Mesh(jax.devices(), axis_names=("batch",)) + replicated = jshard.NamedSharding(mesh, P()) + dp_sharding = jshard.NamedSharding(mesh, P("batch")) + self._dp_sharding = dp_sharding + + batch_sharding = TrainingBatch( + inputs=dp_sharding, + probabilities=dp_sharding, + values=dp_sharding, + ) + in_shardings = (replicated, batch_sharding) + out_shardings = replicated + + jit_kwargs["in_shardings"] = in_shardings + jit_kwargs["out_shardings"] = out_shardings + + @partial(jax.jit, **jit_kwargs) + def _step( + optimizer_tx: optax.GradientTransformation, + jit_state: JitTrainingState, + batch: TrainingBatch, + ) -> Tuple[JitTrainingState, MetricsDict]: + model = nnx.merge(graphdef, jit_state.model_state) + + def loss_for_grad( + model_arg: LczeroModel, sample_arg: TrainingSample + ) -> Tuple[jax.Array, Dict[str, jax.Array]]: + return loss_fn(model_arg, sample_arg) + + loss_vfn = jax.vmap( + loss_for_grad, + in_axes=(None, 0), # (model_arg, sample_arg) + out_axes=0, + ) + + def mean_loss_for_grad( + model_arg: LczeroModel, batch_arg: TrainingBatch + ) -> Tuple[jax.Array, Dict[str, jax.Array]]: + # vmap automatically distributes TrainingBatch over batch dimension, + # calling loss_for_grad with TrainingSample (single samples). + per_sample_data_loss, unweighted_losses = loss_vfn( + model_arg, + batch_arg, # type: ignore[arg-type] + ) + mean_loss = jnp.mean(per_sample_data_loss) + return mean_loss, unweighted_losses + + grad_fn = nnx.value_and_grad(mean_loss_for_grad, has_aux=True) + (mean_loss, unweighted_losses), mean_grads = grad_fn(model, batch) + grad_norm = optax.global_norm(mean_grads) + + assert jit_state.opt_state is not None + updates, new_opt_state = optimizer_tx.update( + mean_grads, jit_state.opt_state, jit_state.model_state + ) + new_model_state = optax.apply_updates( + jit_state.model_state, updates + ) + + new_jit_state = jit_state.replace( + step=jit_state.step + 1, + model_state=new_model_state, + opt_state=new_opt_state, + ) + + mean_unweighted = tree_util.tree_map(jnp.mean, unweighted_losses) + metrics: MetricsDict = { + "loss": mean_loss, + "unweighted_losses": mean_unweighted, + "grad_norm": grad_norm, + } + return new_jit_state, metrics + + self.train_step = cast( + Callable[ + [optax.GradientTransformation, JitTrainingState, TrainingBatch], + Tuple[JitTrainingState, MetricsDict], + ], + _step, + ) + + @staticmethod + @jax.jit + def _swa_tree_map( + alpha: jax.Array, + beta: jax.Array, + swa_state: nnx.State, + model_state: nnx.State, + ) -> nnx.State: + return tree_util.tree_map( + lambda a, b: alpha * a + beta * b, swa_state, model_state + ) + + def update_swa( + self, jit_state: JitTrainingState, weight: float + ) -> JitTrainingState: + """Update SWA using the provided weight for the current model. + + Assumes `jit_state.swa_state` is initialized and `_swa_config` present. + """ + logger.info( + "Updating SWA model, weight=%f, num_averages=%f", + weight, + jit_state.num_averages, + ) + assert self._swa_config is not None + assert jit_state.swa_state is not None + assert weight > 0.0 + max_num_averages = self._swa_config.num_averages + denom = jit_state.num_averages + weight + alpha = jit_state.num_averages / denom + beta = weight / denom + new_swa_state = self._swa_tree_map( + jnp.array(alpha), + jnp.array(beta), + jit_state.swa_state, + jit_state.model_state, + ) + new_num_averages = min( + max_num_averages, jit_state.num_averages + weight + ) + return jit_state.replace( + swa_state=new_swa_state, num_averages=new_num_averages + ) + + def maybe_update_swa( + self, + jit_state: JitTrainingState, + steps_completed: int, + total_steps: int, + ) -> JitTrainingState: + """Optionally update SWA based on configured schedule and epoch progress. + + Returns the original jit_state when no update is scheduled. + """ + if self._swa_config is None: + return jit_state + period_steps = self._swa_config.period_steps + assert period_steps > 0 + if steps_completed % period_steps == 0: + return self.update_swa(jit_state, 1.0) + if steps_completed == total_steps: + remainder = total_steps % period_steps + return self.update_swa(jit_state, remainder / period_steps) + return jit_state + + def _validate_and_prepare_batch( + self, tensor_tuple: tuple[np.ndarray, ...] + ) -> TrainingBatch: + logger.info("Fetched batch from dataloader") + + # Convert tuple to TrainingBatch + batch = TrainingBatch.from_tuple(tensor_tuple) + + # Ensure batch.inputs is jax.Array for shape access + assert isinstance(batch.inputs, jax.Array) + batch_size = batch.inputs.shape[0] + if self._dp_sharding is not None: + num_devices = jax.device_count() + if batch_size % num_devices != 0: + raise ValueError( + f"Batch size {batch_size} must be divisible by device " + f"count {num_devices} for multi-GPU training. " + f"Per-device batch size would be " + f"{batch_size / num_devices:.2f}" + ) + per_device_batch_size = batch_size // num_devices + logger.info( + f"Multi-GPU batch: {batch_size} total " + f"({per_device_batch_size} per device)" + ) + + if self._dp_sharding is not None: + batch = jax.device_put(batch, self._dp_sharding) + + return batch + + def _log_step_metrics( + self, + step_value: int, + local_step: int, + num_steps: int, + metrics: MetricsDict, + ) -> None: + loss = float(metrics["loss"]) + unweighted_losses = { + k: float(v) for k, v in metrics["unweighted_losses"].items() + } + grad_norm = float(metrics["grad_norm"]) + logger.info( + f"Step {step_value} ({local_step}/{num_steps}), Loss: {loss}, " + f"Unweighted losses: {unweighted_losses}, Grad norm: {grad_norm}" + ) + + def _execute_step_hook( + self, + step_hook: Optional[StepHook], + step_value: int, + local_step: int, + num_steps: int, + metrics: MetricsDict, + jit_state: JitTrainingState, + ) -> None: + if step_hook is None: + return + hook_data = StepHookData( + global_step=step_value, + local_step=local_step, + steps_per_epoch=num_steps, + metrics=metrics, + jit_state=jit_state, + ) + step_hook(hook_data) + + def run( + self, + jit_state: JitTrainingState, + datagen: Generator[tuple[np.ndarray, ...], None, None], + num_steps: int, + step_hook: Optional[StepHook] = None, + memory_profile_dir: Optional[str] = None, + ) -> JitTrainingState: + assert jit_state.opt_state is not None + if self._dp_sharding is not None: + replicated = jshard.NamedSharding(self._dp_sharding.mesh, P()) + jit_state = jax.device_put(jit_state, replicated) + for local_step in range(num_steps): + logger.info(f"Starting step {jit_state.step}") + if memory_profile_dir is not None: + jax.profiler.save_device_memory_profile( + f"{memory_profile_dir}/" + f"{datetime.now().strftime('%Y%m%d-%H%M%S')}" + f"_before_{int(jit_state.step)}.prof" + ) + batch = self._validate_and_prepare_batch(next(datagen)) + jit_state, metrics = self.train_step( + self.optimizer_tx, jit_state, batch + ) + step_value = int( + np.asarray(jax.device_get(jit_state.step)).reshape(()) + ) + jit_state = self.maybe_update_swa( + jit_state, local_step + 1, num_steps + ) + self._execute_step_hook( + step_hook, step_value, local_step, num_steps, metrics, jit_state + ) + self._log_step_metrics(step_value, local_step, num_steps, metrics) + return jit_state diff --git a/src/lczero_training/training/tune_lr.py b/src/lczero_training/training/tune_lr.py new file mode 100644 index 00000000..6b8002d6 --- /dev/null +++ b/src/lczero_training/training/tune_lr.py @@ -0,0 +1,282 @@ +import csv +import logging +import sys +from contextlib import nullcontext +from functools import partial +from typing import Callable, Dict, List, Tuple, cast + +import jax +import jax.numpy as jnp +import matplotlib.pyplot as plt +import optax +import orbax.checkpoint as ocp +from flax import nnx +from google.protobuf import text_format +from jax import tree_util + +from lczero_training.dataloader import make_dataloader +from lczero_training.model.loss_function import LczeroLoss +from lczero_training.model.model import LczeroModel +from lczero_training.training.state import TrainingState +from proto.root_config_pb2 import RootConfig + +from .training import Training, from_dataloader + +logger = logging.getLogger(__name__) + + +def _prepare_batch(batch_tuple: tuple) -> Dict: + # DataLoader now returns tuple: (inputs, probabilities, values) + return { + "inputs": batch_tuple[0], + "probabilities": batch_tuple[1], + "values": batch_tuple[2], + } + + +def _make_optimizer_with_schedule( + training_state: TrainingState, + config: RootConfig, + schedule: optax.Schedule, +) -> optax.GradientTransformation: + max_grad_norm = getattr(config.training, "max_grad_norm", 0.0) + opt_config = config.training.optimizer + + if opt_config.HasField("nadamw"): + conf = opt_config.nadamw + tx: optax.GradientTransformation = optax.nadamw( + schedule, + b1=conf.beta_1, + b2=conf.beta_2, + eps=conf.epsilon, + weight_decay=conf.weight_decay, + ) + else: + raise ValueError( + f"Unsupported optimizer type: {opt_config.WhichOneof('optimizer_type')}" + ) + + if max_grad_norm > 0: + tx = optax.chain(optax.clip_by_global_norm(max_grad_norm), tx) + + if training_state.jit_state.opt_state is None: + raise ValueError("Optimizer state must be available in the checkpoint.") + + return tx + + +def _make_eval_step( + graphdef: nnx.GraphDef, loss_fn: LczeroLoss +) -> Callable[[nnx.State, Dict], jax.Array]: + @partial(nnx.jit, static_argnames=()) + def eval_step(model_state: nnx.State, batch: Dict) -> jax.Array: + model = nnx.merge(graphdef, model_state) + + def calculate_loss( + model_arg: LczeroModel, batch_arg: Dict + ) -> Tuple[jax.Array, Dict[str, jax.Array]]: + return loss_fn(model_arg, **batch_arg) + + loss_vfn = jax.vmap(calculate_loss, in_axes=(None, 0), out_axes=0) + per_sample_data_loss, _ = loss_vfn(model, batch) + return jnp.mean(per_sample_data_loss) + + return cast(Callable[[nnx.State, Dict], jax.Array], eval_step) + + +def _plot_results(results: List[Tuple[float, float]], plot_output: str) -> None: + logger.info("Saving plot to %s", plot_output) + lrs, losses = zip(*results) + plt.figure(figsize=(10, 6)) + plt.plot(lrs, losses, marker="o") + plt.xscale("log") + plt.xlabel("Learning Rate") + plt.ylabel("Loss") + plt.title("Learning Rate Finder") + plt.grid(True, which="both", linestyle="--") + plt.tight_layout() + plt.savefig(plot_output, dpi=150) + plt.close() + + +def tune_lr( + *, + config_filename: str, + start_lr: float, + num_steps: int, + multiplier: float = 1.01, + warmup_steps: int = 0, + warmup_lr: float | None = None, + csv_output: str | None = None, + plot_output: str | None = None, + num_test_batches: int = 0, +) -> None: + if num_steps <= 0 or start_lr <= 0 or multiplier <= 0 or warmup_steps < 0: + logger.error( + "num_steps, start_lr, and multiplier must be positive, " + "and warmup_steps non-negative." + ) + sys.exit(1) + if warmup_steps > 0 and (warmup_lr is None or warmup_lr <= 0): + logger.error("warmup_lr must be a positive value when warmup_steps > 0") + sys.exit(1) + + config = RootConfig() + logger.info("Reading configuration from %s", config_filename) + with open(config_filename, "r") as f: + text_format.Parse(f.read(), config) + + if not config.training.checkpoint.path: + logger.error("Checkpoint path must be set in the configuration.") + sys.exit(1) + + checkpoint_mgr = ocp.CheckpointManager( + config.training.checkpoint.path, + options=ocp.CheckpointManagerOptions(create=False), + ) + + logger.info("Creating state from configuration") + empty_state = TrainingState.new_from_config( + model_config=config.model, training_config=config.training + ) + + logger.info("Restoring checkpoint from %s", config.training.checkpoint.path) + training_state = checkpoint_mgr.restore( + checkpoint_mgr.latest_step(), args=ocp.args.PyTreeRestore(empty_state) + ) + if training_state is None: + logger.error("No checkpoint found.") + sys.exit(1) + logger.info("Restored checkpoint at step %d", training_state.jit_state.step) + + assert isinstance(training_state, TrainingState) + + model, _ = nnx.split( + LczeroModel(config=config.model, rngs=nnx.Rngs(params=42)) + ) + + datagen = from_dataloader(make_dataloader(config.data_loader)) + + # Prepare fixed validation batches only if requested (num_test_batches > 0). + use_validation = num_test_batches > 0 + if use_validation: + logger.info("Fetching %d validation batches", num_test_batches) + validation_batches = [ + tree_util.tree_map(jnp.asarray, _prepare_batch(next(datagen))) + for _ in range(num_test_batches) + ] + else: + validation_batches = [] + + loss_fn = LczeroLoss(config=config.training.losses) + eval_step = _make_eval_step(model, loss_fn) + + def avg_val_loss() -> float: + assert use_validation + total_loss = 0.0 + for vb in validation_batches: + total_loss += float( + eval_step(training_state.jit_state.model_state, vb) + ) + return total_loss / float(num_test_batches) + + def train_one_step( + training: Training, tx: optax.GradientTransformation + ) -> float: + nonlocal training_state + batch = tree_util.tree_map(jnp.asarray, _prepare_batch(next(datagen))) + new_jit_state, metrics = training.train_step( + tx, training_state.jit_state, batch + ) + training_state = training_state.replace(jit_state=new_jit_state) + return float(metrics["loss"]) # training batch loss + + def run_phase( + *, + steps: int, + schedule: optax.Schedule, + lr_at: Callable[[int], float], + label: str, + on_result: Callable[[float, float, float | None], None], + ) -> None: + start_step = training_state.jit_state.step + + def offset_schedule(count: jax.Array) -> jax.Array: + return schedule(count - start_step) + + tx = _make_optimizer_with_schedule( + training_state, config, offset_schedule + ) + training = Training(optimizer_tx=tx, graphdef=model, loss_fn=loss_fn) + for i in range(steps): + current_lr = lr_at(i) + logger.info( + "%s step %d/%d at lr %.8f", label, i + 1, steps, current_lr + ) + train_loss = train_one_step(training, tx) + val_loss = avg_val_loss() if use_validation else None + on_result(current_lr, train_loss, val_loss) + if use_validation: + logger.info( + "%s at lr %.8f: train=%.6f, val=%.6f", + label, + current_lr, + train_loss, + cast(float, val_loss), + ) + else: + logger.info( + "%s train loss at lr %.8f: %.6f", + label, + current_lr, + train_loss, + ) + + results: List[Tuple[float, float]] = [] + with ( + open(csv_output, "w", newline="") if csv_output else nullcontext() + ) as csv_file: + writer = csv.writer(csv_file) if csv_file else None + if writer: + if use_validation: + writer.writerow(["lr", "train_loss", "val_loss"]) + else: + writer.writerow(["lr", "train_loss"]) + + def on_result( + lr: float, train_loss: float, val_loss: float | None + ) -> None: + results.append((lr, train_loss)) + if writer and csv_file: + if use_validation: + writer.writerow([lr, train_loss, val_loss]) + else: + writer.writerow([lr, train_loss]) + csv_file.flush() + + phases = [] + if warmup_steps > 0 and warmup_lr is not None: + phases.append( + { + "label": "Warmup", + "steps": warmup_steps, + "schedule": optax.constant_schedule(warmup_lr), + "lr_at": lambda _: float(warmup_lr), + } + ) + phases.append( + { + "label": "Sweep", + "steps": num_steps, + "schedule": optax.exponential_decay( + start_lr, transition_steps=1, decay_rate=multiplier + ), + "lr_at": lambda i: start_lr * (multiplier**i), + } + ) + + for phase_params in phases: + run_phase(**phase_params, on_result=on_result) + + if plot_output: + _plot_results(results, plot_output) diff --git a/src/lczero_training/training/utils.py b/src/lczero_training/training/utils.py new file mode 100644 index 00000000..30e3e363 --- /dev/null +++ b/src/lczero_training/training/utils.py @@ -0,0 +1,20 @@ +from pathlib import PurePosixPath + +from flax import nnx + +from proto.training_config_pb2 import WeightsSelector + + +def make_weights_mask( + selector: WeightsSelector, params: nnx.State +) -> nnx.State: + """Creates a boolean mask based on WeightsSelector. True = include weight.""" + + def mask_fn(path: tuple[object, ...], _variable: nnx.Variable) -> bool: + p = PurePosixPath(*map(str, path)) + for rule in selector.rule: + if p.full_match(rule.match): + return rule.include + return selector.otherwise_include + + return nnx.map_state(mask_fn, params) diff --git a/src/lczero_training/tui/__init__.py b/src/lczero_training/tui/__init__.py new file mode 100644 index 00000000..15915b96 --- /dev/null +++ b/src/lczero_training/tui/__init__.py @@ -0,0 +1,6 @@ +# ABOUTME: TUI package initialization for the training dashboard. +# ABOUTME: Exports main TrainingTuiApp class for external use. + +from .app import TrainingTuiApp + +__all__ = ["TrainingTuiApp"] diff --git a/src/lczero_training/tui/app.py b/src/lczero_training/tui/app.py new file mode 100644 index 00000000..88f7d771 --- /dev/null +++ b/src/lczero_training/tui/app.py @@ -0,0 +1,244 @@ +# ABOUTME: Main TUI application class implementing the training dashboard. +# ABOUTME: Uses Textual framework to create a full-screen interface with four panes. + +import argparse +import logging +import os +import signal +import subprocess +import sys +from datetime import datetime +from typing import Iterable, Optional + +import anyio +from anyio.streams.text import TextReceiveStream, TextSendStream +from textual.app import App, ComposeResult, SystemCommand +from textual.containers import Horizontal +from textual.screen import Screen +from textual.widgets import Footer, Static + +from ..daemon.protocol.communicator import AsyncCommunicator +from ..daemon.protocol.messages import ( + StartTrainingImmediatelyPayload, + StartTrainingPayload, + TrainingStatusPayload, +) +from .data_pipeline_pane import DataPipelinePane +from .log_pane import StreamingLogPane +from .training_widgets import TrainingScheduleWidget + +logger = logging.getLogger(__name__) + + +class HeaderBar(Static): + """Empty header bar.""" + + def compose(self) -> ComposeResult: + return + yield # unreachable, but makes the function a generator + + +class JAXTrainingPane(Static): + """Right pane showing JAX training status and metrics.""" + + def compose(self) -> ComposeResult: + yield Static( + "JAX Training Status\n\n" + "Live training metrics will be displayed here when active:\n" + "• Epoch Progress\n• Performance Metrics\n• Loss Values", + classes="jax-training-content", + ) + + +class TrainingTuiApp(App): + """Main TUI application for the training dashboard. + + This creates a full-screen interface with four main panes: + - Header bar with uptime and status + - Data pipeline pane (main/left) + - Training status pane (right) + - Log pane (bottom) + """ + + CSS_PATH = "app.tcss" + + @staticmethod + def add_arguments(parser: argparse.ArgumentParser) -> None: + """Adds all required command-line arguments to the given parser.""" + parser.add_argument( + "--config", + required=True, + help="Path to the training configuration file", + ) + parser.add_argument( + "--logfile", + help="Path to the log file for saving TUI output", + ) + parser.add_argument( + "--io-dump", + help="Path to file for dumping raw daemon IO for debugging", + ) + parser.add_argument( + "--daemon-flag", + action="append", + default=[], + dest="daemon_flags", + help="Extra argument to pass to the daemon (repeatable)", + ) + + _log_stream: TextReceiveStream + _daemon_process: anyio.abc.Process + _communicator: AsyncCommunicator + _config_file: str + _logfile: Optional[str] + _data_pipeline_pane: DataPipelinePane + _training_schedule_widget: TrainingScheduleWidget + + BINDINGS = [ + ("q", "quit", "Quit"), + ("ctrl+c", "quit", "Quit"), + ] + + def __init__(self, args: Optional[argparse.Namespace] = None) -> None: + """ + Initializes the app. + If 'args' is provided, it's used directly. + If 'args' is None, fallback to parsing sys.argv. + """ + super().__init__() + + if args is None: + # Fallback for when run by "textual run" + parser = argparse.ArgumentParser() + TrainingTuiApp.add_arguments(parser) + args, _ = parser.parse_known_args() + + # Consume configuration from the args object + self._config_file: str = args.config + self._logfile: Optional[str] = args.logfile + self._io_dump_file: Optional[str] = args.io_dump + self._daemon_flags: list[str] = args.daemon_flags + + async def on_load(self) -> None: + """Start the daemon process and communicator when the app loads.""" + # Create the daemon process via Python module execution to avoid PATH reliance. + env = None + if "TF_CPP_MIN_LOG_LEVEL" not in os.environ: + env = {**os.environ, "TF_CPP_MIN_LOG_LEVEL": "0"} + + self._daemon_process = await anyio.open_process( + [ + sys.executable, + "-m", + "lczero_training.commands.daemon", + *self._daemon_flags, + ], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + ) + + assert self._daemon_process.stderr is not None + assert self._daemon_process.stdin is not None + assert self._daemon_process.stdout is not None + + # Set up streams and communicator + self._log_stream = TextReceiveStream(self._daemon_process.stderr) + io_dump = None + if self._io_dump_file: + io_dump = open(self._io_dump_file, "a", buffering=1) + ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + io_dump.write(f"======= {ts} =======\n") + self._communicator = AsyncCommunicator( + handler=self, + input_stream=TextReceiveStream(self._daemon_process.stdout), + output_stream=TextSendStream(self._daemon_process.stdin), + io_dump=io_dump, + ) + + def compose(self) -> ComposeResult: + """Compose the main UI layout.""" + yield HeaderBar() + + self._data_pipeline_pane = DataPipelinePane() + self._data_pipeline_pane.border_title = "Training data pipeline" + yield self._data_pipeline_pane + + # Horizontal split below the data pipeline pane + with Horizontal(id="training-status-container"): + self._training_schedule_widget = TrainingScheduleWidget() + self._training_schedule_widget.border_title = "Training Schedule" + yield self._training_schedule_widget + + jax_training_pane = JAXTrainingPane() + jax_training_pane.border_title = "JAX Training Status" + yield jax_training_pane + + yield StreamingLogPane( + stream=self._log_stream, logfile_path=self._logfile + ) + yield Footer() + + async def _monitor_daemon_process(self) -> None: + """Notify and log when the daemon process exits unexpectedly.""" + await self._daemon_process.wait() + rc = self._daemon_process.returncode + sig = -rc if rc is not None and rc < 0 else None + if sig is not None: + msg = f"Daemon killed by signal {sig} ({signal.Signals(sig).name})" + else: + msg = f"Daemon exited with code {rc}" + logger.warning(msg) + self.notify(msg, severity="warning", timeout=60) + if self._logfile: + with open(self._logfile, "a") as f: + f.write(f"{msg}\n") + + def on_mount(self) -> None: + """Start the communicator when the app mounts.""" + self.run_worker(self._communicator.run(), exclusive=True) + self.run_worker(self._send_start_training(), exclusive=False) + self.run_worker(self._monitor_daemon_process(), exclusive=False) + + async def _send_start_training(self) -> None: + """Send StartTrainingPayload with the config file.""" + payload = StartTrainingPayload(config_filepath=self._config_file) + await self._communicator.send(payload) + + async def _command_start_training_immediately(self) -> None: + """Trigger immediate training without waiting for additional chunks.""" + + payload = StartTrainingImmediatelyPayload() + await self._communicator.send(payload) + self.notify("Requested immediate training start.") + + async def action_quit(self) -> None: # type: ignore + """Handle quit action.""" + self._daemon_process.send_signal(signal.SIGINT) + with anyio.move_on_after(20) as scope: + await self._daemon_process.wait() + if scope.cancelled_caught: + self._daemon_process.terminate() + self.exit() + + def get_system_commands(self, screen: Screen) -> Iterable[SystemCommand]: + """Add application-specific commands to the command palette.""" + + yield from super().get_system_commands(screen) + yield SystemCommand( + "Start training immediately", + "Begin the next training cycle without waiting for chunks.", + self._command_start_training_immediately, + ) + + async def on_training_status(self, payload: TrainingStatusPayload) -> None: + """Handle training status updates.""" + self._data_pipeline_pane.update_metrics( + payload.dataloader_1_second, payload.dataloader_total + ) + + # Update training schedule widget + self._training_schedule_widget.update_training_schedule( + payload.training_schedule + ) diff --git a/src/lczero_training/tui/app.tcss b/src/lczero_training/tui/app.tcss new file mode 100644 index 00000000..ad68c9a3 --- /dev/null +++ b/src/lczero_training/tui/app.tcss @@ -0,0 +1,236 @@ +Screen { + background: $primary; +} + +HeaderBar { + dock: top; + height: 1; + background: $primary-darken-1; + color: $text; +} + +DataPipelinePane { + background: $surface; + color: $text; + border: none; + margin: 1; + padding: 0 1; + layout: vertical; + overflow-y: auto; + height: 27; +} + +#training-status-container { + height: 12; + layout: horizontal; +} + +TrainingScheduleWidget { + background: $surface; + color: $text; + border: solid $primary; + margin: 1; + width: 1fr; +} + +JAXTrainingPane { + background: $surface; + color: $text; + border: solid $primary; + margin: 1; + width: 1fr; +} + +RichLog { + height: 1fr; + background: $panel; + color: $text; + margin: 0; + width: 100%; + overflow-y: scroll; +} + +.stage-content, .queue-content { + text-align: left; +} + +.dataloader-row { + layout: grid; + grid-size: 6 8; + grid-columns: 1fr 1fr 1fr 1fr 1fr 1fr; + grid-rows: 1; + grid-gutter: 0; + border: none; + padding: 0; + margin: 0; + height: auto; + min-height: 1; + width: 100%; +} + +.stage-row { + background: $surface; + color: $text; + grid-size: 5 8; + grid-columns: 1fr 1fr 1fr 1fr 1fr; +} + +.queue-row { + background: $surface-darken-2; + color: $text; + grid-size: 6 1; +} + +.statistics-row { + background: $surface-darken-1; + color: $text; + grid-size: 1 1; +} + +.statistics-full { + background: $surface-darken-1; + color: $text; + padding: 0 1; + margin: 0; + height: auto; + min-height: 1; + width: 100%; +} + +.row-label { + padding: 0 1 0 0; + margin: 0; + height: 1; + max-height: 1; + text-style: bold; + color: $text; + content-align: left middle; + text-wrap: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} + +.grid-spacer { + width: 1; + height: 1; + max-height: 1; + background: transparent; +} + +.metric-chip { + background: $surface-darken-1; + color: $text; + padding: 0 1; + margin: 0 1 0 0; + height: 1; + max-height: 1; + border: none; + content-align: center middle; + width: auto; + min-width: 10; + text-wrap: nowrap; + overflow: hidden; +} + +.load-chip { + background: $primary-darken-2; + color: $text; +} + +.info-chip { + background: $surface-darken-1; + color: $text; +} + +.warning-chip { + background: $warning; + color: $text; +} + +.queue-name-chip { + background: transparent; + color: $text-muted; + padding: 0; + margin-right: 1; + height: 1; + max-height: 1; + text-wrap: nowrap; + overflow: hidden; +} + +.queue-fill { + height: 1; + max-height: 1; + width: 1fr; + margin-right: 1; +} + +.queue-fill-text { + background: transparent; + color: $text; + margin: 0; + padding: 0; + height: 1; + max-height: 1; + text-wrap: nowrap; + overflow: hidden; + width: auto; +} + +.queue-rate--zero { + color: $error; +} + +.queue-rate, .queue-total { + text-wrap: nowrap; +} + + +.jax-training-content { + padding: 1; +} + +/* Training widgets styles */ +.time-label { + width: auto; + padding: 0; + margin-right: 1; + text-align: left; +} + +.time-progress-bar { + width: 1fr; + height: 1; +} + +.time-progress-bar .bar--indeterminate { + background: $panel; +} + +.time-progress-bar .bar--bar { + background: $panel; +} + +.time-ratio { + width: auto; + padding: 0; + text-align: right; +} + +TimeProgressWidget { + height: 1; + layout: horizontal; +} + +ChunksProgressWidget { + height: 1; + layout: horizontal; +} + +#uptime-stage-display, #epochs-display { + height: 1; +} + +.log-content { + padding: 0; +} diff --git a/src/lczero_training/tui/data_pipeline_pane.py b/src/lczero_training/tui/data_pipeline_pane.py new file mode 100644 index 00000000..8b884108 --- /dev/null +++ b/src/lczero_training/tui/data_pipeline_pane.py @@ -0,0 +1,189 @@ +# ABOUTME: Data pipeline pane widget for displaying DataLoader metrics. +# ABOUTME: Shows a grid of pipeline stages and queues with their metrics. + +from collections.abc import Iterable +from typing import Any + +from textual.app import ComposeResult +from textual.containers import Container + +import proto.training_metrics_pb2 as training_metrics_pb2 + +from .dataloader_widgets import QueueWidget, StageWidget, StatisticsRowWidget + +FRIENDLY_STAGE_NAMES = { + "file_path_provider": "File discovery", + "chunk_source_loader": "Chunk source loader", + "shuffling_chunk_pool": "Shuffling chunk pool", + "chunk_rescorer": "Chunk rescorer", + "chunk_splitter": "Chunk splitter", + "chunk_unpacker": "Chunk unpacker", + "shuffling_frame_sampler": "Shuffling frame sampler", + "tensor_generator": "Batched tensor generator", +} + + +class DataPipelinePane(Container): + """Main pane showing data pipeline flow and statistics as a grid.""" + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._stage_widgets: dict[str, StageWidget] = {} + self._queue_widgets: dict[str, dict[str, QueueWidget]] = {} + self._queue_order: dict[str, list[str]] = {} + self._statistics_widgets: dict[str, dict[str, StatisticsRowWidget]] = {} + self._statistics_order: dict[str, list[str]] = {} + self._stage_order: list[str] = [] + + def compose(self) -> ComposeResult: + """The pane starts empty and rows are added when metrics arrive.""" + yield from () + + def _friendly_title(self, stage_key: str) -> str: + return FRIENDLY_STAGE_NAMES.get( + stage_key, stage_key.replace("_", " ").title() + ) + + def _ensure_stage_widget( + self, + stage_key: str, + ) -> tuple[StageWidget, bool]: + created = False + if stage_key not in self._stage_widgets: + stage_widget = StageWidget( + stage_key=stage_key, + fallback_name=self._friendly_title(stage_key), + ) + self._stage_widgets[stage_key] = stage_widget + self._queue_widgets[stage_key] = {} + self._queue_order[stage_key] = [] + self._stage_order.append(stage_key) + created = True + else: + stage_widget = self._stage_widgets[stage_key] + + return stage_widget, created + + def _ensure_queue_widgets( + self, + stage_key: str, + stage_metric: training_metrics_pb2.StageMetricProto, + ) -> list[QueueWidget]: + stage_queue_widgets = self._queue_widgets.setdefault(stage_key, {}) + queue_order = self._queue_order.setdefault(stage_key, []) + new_widgets: list[QueueWidget] = [] + + friendly = self._friendly_title(stage_key) + for index, queue_metric in enumerate(stage_metric.queue_metrics): + queue_identifier = queue_metric.name or f"__index__{index}" + if queue_identifier in stage_queue_widgets: + continue + label_suffix = ( + f" {queue_metric.name}" + if queue_metric.name + else f" #{index + 1}" + ) + queue_widget = QueueWidget( + stage_key=stage_key, + stage_name=f"{friendly} queue{label_suffix}", + queue_name=queue_identifier, + ) + stage_queue_widgets[queue_identifier] = queue_widget + queue_order.append(queue_identifier) + new_widgets.append(queue_widget) + + return new_widgets + + def _ensure_statistics_widgets( + self, + stage_key: str, + stage_metric: training_metrics_pb2.StageMetricProto, + ) -> list[StatisticsRowWidget]: + stage_statistics_widgets = self._statistics_widgets.setdefault( + stage_key, {} + ) + statistics_order = self._statistics_order.setdefault(stage_key, []) + new_widgets: list[StatisticsRowWidget] = [] + + for statistics_metric in stage_metric.statistics_metrics: + metric_name = statistics_metric.name or "" + if metric_name in stage_statistics_widgets: + continue + label = metric_name.replace("_", " ").title() + statistics_widget = StatisticsRowWidget( + stage_key=stage_key, + metric_name=metric_name, + label=label, + ) + stage_statistics_widgets[metric_name] = statistics_widget + statistics_order.append(metric_name) + new_widgets.append(statistics_widget) + + return new_widgets + + def _mount_widgets( + self, widgets: Iterable[StageWidget | QueueWidget | StatisticsRowWidget] + ) -> None: + async def _do_mount() -> None: + await self.mount(*widgets) + + self.call_later(_do_mount) + + def _ensure_rows( + self, + metrics: training_metrics_pb2.DataLoaderMetricsProto, + ) -> None: + new_widgets: list[StageWidget | QueueWidget | StatisticsRowWidget] = [] + for stage_metric in metrics.stage_metrics: + stage_key = stage_metric.name + if not stage_key: + continue + + stage_widget, created = self._ensure_stage_widget(stage_key) + if created: + new_widgets.append(stage_widget) + + statistics_widgets = self._ensure_statistics_widgets( + stage_key, stage_metric + ) + new_widgets.extend(statistics_widgets) + + queue_widgets = self._ensure_queue_widgets(stage_key, stage_metric) + new_widgets.extend(queue_widgets) + + if new_widgets: + self._mount_widgets(new_widgets) + + def update_metrics( + self, + dataloader_1_second: training_metrics_pb2.DataLoaderMetricsProto | None, + dataloader_total: training_metrics_pb2.DataLoaderMetricsProto | None, + ) -> None: + """Update all pipeline stages and queues with new metrics.""" + + metrics_for_layout = dataloader_total or dataloader_1_second + if metrics_for_layout: + self._ensure_rows(metrics_for_layout) + + for stage_key in self._stage_order: + stage_widget = self._stage_widgets.get(stage_key) + if stage_widget: + stage_widget.update_metrics( + dataloader_1_second, dataloader_total + ) + + for statistics_key in self._statistics_order.get(stage_key, []): + statistics_widget = self._statistics_widgets[stage_key].get( + statistics_key + ) + if statistics_widget: + statistics_widget.update_metrics( + dataloader_1_second, dataloader_total + ) + + for queue_key in self._queue_order.get(stage_key, []): + queue_widget = self._queue_widgets[stage_key].get(queue_key) + if queue_widget: + queue_widget.update_metrics( + dataloader_1_second, dataloader_total + ) diff --git a/src/lczero_training/tui/dataloader_widgets.py b/src/lczero_training/tui/dataloader_widgets.py new file mode 100644 index 00000000..f97707c5 --- /dev/null +++ b/src/lczero_training/tui/dataloader_widgets.py @@ -0,0 +1,593 @@ +"""Widgets that render data loader metrics without stage-specific logic.""" + +from __future__ import annotations + +from typing import Any, Dict + +from textual.app import ComposeResult +from textual.widget import Widget +from textual.widgets import ProgressBar, Static + +import proto.training_metrics_pb2 as training_metrics_pb2 + + +def _find_stage_metric( + metrics: training_metrics_pb2.DataLoaderMetricsProto | None, + stage_key: str, +) -> training_metrics_pb2.StageMetricProto | None: + if not metrics: + return None + for stage_metric in metrics.stage_metrics: + if stage_metric.name == stage_key: + return stage_metric + return None + + +def _collect_metric_names( + stage_1s: training_metrics_pb2.StageMetricProto | None, + stage_total: training_metrics_pb2.StageMetricProto | None, + attribute: str, +) -> list[str]: + names: list[str] = [] + + def _add_from( + stage_metric: training_metrics_pb2.StageMetricProto | None, + ) -> None: + if not stage_metric: + return + for metric in getattr(stage_metric, attribute): + name = metric.name if metric.name else "" + if name not in names: + names.append(name) + + _add_from(stage_total) + _add_from(stage_1s) + return names + + +def _find_load_metric( + stage_metric: training_metrics_pb2.StageMetricProto | None, + metric_name: str, +) -> training_metrics_pb2.LoadMetricProto | None: + if not stage_metric: + return None + for load_metric in stage_metric.load_metrics: + if (load_metric.name or "") == metric_name: + return load_metric + return None + + +def _find_count_metric( + stage_metric: training_metrics_pb2.StageMetricProto | None, + metric_name: str, +) -> training_metrics_pb2.CountMetricProto | None: + if not stage_metric: + return None + for count_metric in stage_metric.count_metrics: + if (count_metric.name or "") == metric_name: + return count_metric + return None + + +def _find_gauge_metric( + stage_metric: training_metrics_pb2.StageMetricProto | None, + metric_name: str, +) -> training_metrics_pb2.GaugeMetricProto | None: + if not stage_metric: + return None + for gauge_metric in stage_metric.gauge_metrics: + if (gauge_metric.name or "") == metric_name: + return gauge_metric + return None + + +def _find_statistics_metric( + stage_metric: training_metrics_pb2.StageMetricProto | None, + metric_name: str, +) -> training_metrics_pb2.StatisticsProtoDouble | None: + if not stage_metric: + return None + for stats_metric in stage_metric.statistics_metrics: + if (stats_metric.name or "") == metric_name: + return stats_metric + return None + + +def _get_queue_metric( + stage_metric: training_metrics_pb2.StageMetricProto | None, + queue_name: str | None, +) -> training_metrics_pb2.QueueMetricProto | None: + if not stage_metric or not stage_metric.queue_metrics: + return None + if queue_name is None: + return stage_metric.queue_metrics[0] + if queue_name.startswith("__index__"): + try: + index = int(queue_name.removeprefix("__index__")) + except ValueError: + index = -1 + if 0 <= index < len(stage_metric.queue_metrics): + return stage_metric.queue_metrics[index] + for queue_metric in stage_metric.queue_metrics: + if (queue_metric.name or "") == queue_name: + return queue_metric + return None + + +def format_si(value: int, precision: int = 1) -> str: + if value == 0: + return "0" + units = [ + (1_000_000_000_000, "T"), + (1_000_000_000, "G"), + (1_000_000, "M"), + (1_000, "k"), + ] + for threshold, unit in units: + if value >= threshold: + result = value / threshold + if precision == 0: + return f"{int(result)}{unit}" + return f"{result:.{precision}f}{unit}".rstrip("0").rstrip(".") + return str(value) + + +def format_full_number(value: int) -> str: + if value < 10_000: + return str(value) + return f"{value:_}".replace("_", "'") + + +def _format_load( + load_metric: training_metrics_pb2.LoadMetricProto | None, + label: str, +) -> str: + if not load_metric: + return f"{label} --" + total_part = ( + f"{load_metric.total_seconds:.0f}" + if load_metric.total_seconds > 0 + else "--" + ) + return f"{label} {load_metric.load_seconds:.1f}/{total_part}s" + + +def _format_count( + count_metric_1s: training_metrics_pb2.CountMetricProto | None, + count_metric_total: training_metrics_pb2.CountMetricProto | None, + label: str, +) -> str: + if count_metric_1s and count_metric_total: + rate = format_si(count_metric_1s.count) + total = format_full_number(count_metric_total.count) + return f"{label} {rate}/s ({total} total)" + if count_metric_total: + return f"{label} {format_full_number(count_metric_total.count)}" + if count_metric_1s: + return f"{label} {format_si(count_metric_1s.count)}/s" + return f"{label} --" + + +def _format_gauge( + gauge_metric: training_metrics_pb2.GaugeMetricProto | None, + label: str, +) -> str: + if not gauge_metric: + return f"{label} --" + if gauge_metric.HasField("capacity"): + value_text = format_full_number(gauge_metric.value) + capacity_text = format_full_number(gauge_metric.capacity) + return f"{label} {value_text}/{capacity_text}" + return f"{label} {format_full_number(gauge_metric.value)}" + + +def _format_statistics( + stats_1s: training_metrics_pb2.StatisticsProtoDouble | None, + stats_total: training_metrics_pb2.StatisticsProtoDouble | None, + label: str, +) -> str: + parts = [label] + if stats_1s and stats_1s.count > 0: + avg_1s = stats_1s.sum / stats_1s.count + parts.append( + f"per 1s: avg {avg_1s:.1f} (min {stats_1s.min:.1f}, max {stats_1s.max:.1f}, count {format_si(stats_1s.count)})" + ) + if stats_total and stats_total.count > 0: + avg_total = stats_total.sum / stats_total.count + parts.append( + f"total: avg {avg_total:.1f} (min {stats_total.min:.1f}, max {stats_total.max:.1f}, count {format_full_number(stats_total.count)})" + ) + if len(parts) == 1: + return f"{label} --" + return "; ".join(parts) + + +def _average_queue_fullness( + queue_metric: training_metrics_pb2.QueueMetricProto | None, +) -> int | None: + if not queue_metric: + return None + if ( + queue_metric.HasField("queue_fullness") + and queue_metric.queue_fullness.count > 0 + ): + return int( + queue_metric.queue_fullness.sum / queue_metric.queue_fullness.count + ) + return None + + +def _canonical_stage_name( + stage_metric: training_metrics_pb2.StageMetricProto | None, + fallback: str | None, + stage_key: str | None, +) -> str: + if stage_metric and stage_metric.name: + return stage_metric.name + if fallback: + return fallback + if stage_key: + return stage_key + return "--" + + +class BaseRowWidget(Widget): + """Base class for pipeline rows that renders a name and content widgets.""" + + MAX_GRID_ROWS = 8 # Supports up to 28 chips (4 per row, 7 content rows) + + def __init__( + self, + stage_key: str | None = None, + fallback_name: str | None = None, + row_type: str = "stage-row", + **kwargs: Any, + ) -> None: + classes = f"dataloader-row {row_type}" + super().__init__(classes=classes, **kwargs) + self.stage_key = stage_key + self._fallback_name = fallback_name + self._name_label = Static( + _canonical_stage_name(None, fallback_name, stage_key), + classes="row-label", + ) + self._content_widgets: list[Widget] = [] + self._spacers: list[Static] = [] + + def compose(self) -> ComposeResult: + # Name label goes in first cell (row 0, col 0) + yield self._name_label + + def on_mount(self) -> None: + if self._content_widgets: + + async def _mount_initial() -> None: + # Mount chips with spacers interleaved at the start of each row + for i, widget in enumerate(self._content_widgets): + # After every 4 chips, add a spacer for the next row's column 0 + if i > 0 and i % 4 == 0: + spacer = Static("", classes="grid-spacer") + self._spacers.append(spacer) + await self.mount(spacer) + await self.mount(widget) + + self.call_later(_mount_initial) + + def add_content_widget(self, widget: Widget) -> None: + if widget in self._content_widgets: + return + self._content_widgets.append(widget) + + def _update_name( + self, + stage_metric: training_metrics_pb2.StageMetricProto | None, + ) -> None: + self._name_label.update( + _canonical_stage_name( + stage_metric, self._fallback_name, self.stage_key + ) + ) + + +class StageWidget(BaseRowWidget): + """Row widget that renders all metrics exposed by a stage.""" + + def __init__( + self, + stage_key: str, + fallback_name: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__( + stage_key=stage_key, + fallback_name=fallback_name, + row_type="stage-row", + **kwargs, + ) + self._chips: Dict[str, Static] = {} + + def _ensure_chip(self, key: str, default_text: str, classes: str) -> Static: + chip = self._chips.get(key) + if chip is None: + chip = Static(default_text, classes=f"metric-chip {classes}") + self._chips[key] = chip + self.add_content_widget(chip) + return chip + + def _update_last_chunk_chip( + self, + stage_1s: training_metrics_pb2.StageMetricProto | None, + stage_total: training_metrics_pb2.StageMetricProto | None, + ) -> None: + stage = None + if stage_1s and stage_1s.HasField("last_chunk_key"): + stage = stage_1s + elif stage_total and stage_total.HasField("last_chunk_key"): + stage = stage_total + if not stage: + return + last_value = stage.last_chunk_key or "--" + chip = self._ensure_chip("info:last", "last --", "info-chip") + chip.update(f"last {last_value}") + + def _update_anchor_chip( + self, + stage_total: training_metrics_pb2.StageMetricProto | None, + ) -> None: + if not stage_total or not stage_total.HasField("anchor"): + return + chip = self._ensure_chip("info:anchor", "anchor --", "info-chip") + chip.update(f"anchor {stage_total.anchor}") + + def update_metrics( + self, + dataloader_1_second: training_metrics_pb2.DataLoaderMetricsProto | None, + dataloader_total: training_metrics_pb2.DataLoaderMetricsProto | None, + ) -> None: + if self.stage_key is None: + return + stage_metric_1s = _find_stage_metric( + dataloader_1_second, self.stage_key + ) + stage_metric_total = _find_stage_metric( + dataloader_total, self.stage_key + ) + + self._update_name(stage_metric_1s or stage_metric_total) + + load_names = _collect_metric_names( + stage_metric_1s, stage_metric_total, "load_metrics" + ) + for load_name in load_names: + label = load_name or "load" + load_metric = _find_load_metric(stage_metric_1s, load_name) + if load_metric is None: + load_metric = _find_load_metric(stage_metric_total, load_name) + chip = self._ensure_chip( + f"load:{load_name}", f"{label} --", "load-chip" + ) + chip.update(_format_load(load_metric, label=label)) + + count_names = _collect_metric_names( + stage_metric_1s, stage_metric_total, "count_metrics" + ) + for count_name in count_names: + label = count_name or "count" + count_metric_1s = _find_count_metric(stage_metric_1s, count_name) + count_metric_total = _find_count_metric( + stage_metric_total, count_name + ) + chip = self._ensure_chip( + f"count:{count_name}", f"{label} --", "info-chip" + ) + chip.update( + _format_count(count_metric_1s, count_metric_total, label=label) + ) + + gauge_names = _collect_metric_names( + stage_metric_1s, stage_metric_total, "gauge_metrics" + ) + for gauge_name in gauge_names: + label = gauge_name or "gauge" + gauge_metric = _find_gauge_metric(stage_metric_total, gauge_name) + if gauge_metric is None: + gauge_metric = _find_gauge_metric(stage_metric_1s, gauge_name) + chip = self._ensure_chip( + f"gauge:{gauge_name}", f"{label} --", "info-chip" + ) + chip.update(_format_gauge(gauge_metric, label=label)) + + self._update_last_chunk_chip(stage_metric_1s, stage_metric_total) + self._update_anchor_chip(stage_metric_total) + + +class StatisticsRowWidget(BaseRowWidget): + """Full-width row for statistics metrics.""" + + def __init__( + self, + stage_key: str, + metric_name: str, + label: str, + **kwargs: Any, + ) -> None: + super().__init__( + stage_key=stage_key, + fallback_name=None, + row_type="statistics-row", + **kwargs, + ) + self._metric_name = metric_name + self._label = label + self._stats_label = Static(f"{label} --", classes="statistics-full") + self.add_content_widget(self._stats_label) + + def compose(self) -> ComposeResult: + return + yield + + def on_mount(self) -> None: + if self._content_widgets: + + async def _mount_initial() -> None: + await self.mount(*self._content_widgets) + + self.call_later(_mount_initial) + + def update_metrics( + self, + dataloader_1_second: training_metrics_pb2.DataLoaderMetricsProto | None, + dataloader_total: training_metrics_pb2.DataLoaderMetricsProto | None, + ) -> None: + if not self.stage_key: + self._stats_label.update(f"{self._label} --") + return + stage_1s = _find_stage_metric(dataloader_1_second, self.stage_key) + stage_total = _find_stage_metric(dataloader_total, self.stage_key) + stats_1s = _find_statistics_metric(stage_1s, self._metric_name) + stats_total = _find_statistics_metric(stage_total, self._metric_name) + self._stats_label.update( + _format_statistics(stats_1s, stats_total, self._label) + ) + + +class QueueWidget(BaseRowWidget): + """Row widget for queue metrics between stages.""" + + def __init__( + self, + stage_key: str | None = None, + stage_name: str | None = None, + queue_name: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__( + stage_key=stage_key, + fallback_name=stage_name, + row_type="queue-row", + **kwargs, + ) + self._queue_name = queue_name + self._queue_name_chip = Static( + "queue --", classes="metric-chip queue-name-chip" + ) + self._rate_chip = Static("rate --/s", classes="metric-chip queue-rate") + self._total_chip = Static("total --", classes="metric-chip queue-total") + self._drop_chip = Static("dropped --", classes="metric-chip") + self._fill_bar = ProgressBar( + classes="queue-fill", + show_percentage=False, + show_eta=False, + ) + self._fill_text = Static("--/--", classes="metric-chip queue-fill-text") + self.add_content_widget(self._queue_name_chip) + self.add_content_widget(self._rate_chip) + self.add_content_widget(self._total_chip) + self.add_content_widget(self._drop_chip) + self.add_content_widget(self._fill_bar) + self.add_content_widget(self._fill_text) + + def compose(self) -> ComposeResult: + # Queue widgets don't show the name label - use all 6 columns + return + yield # Make this a generator + + def on_mount(self) -> None: + # Mount all content widgets without spacers (use all 6 columns per row) + if self._content_widgets: + + async def _mount_initial() -> None: + await self.mount(*self._content_widgets) + + self.call_later(_mount_initial) + + def update_metrics( + self, + dataloader_1_second: training_metrics_pb2.DataLoaderMetricsProto | None, + dataloader_total: training_metrics_pb2.DataLoaderMetricsProto | None, + ) -> None: + if not self.stage_key: + self._queue_name_chip.update("queue --") + self._rate_chip.update("rate --/s") + self._total_chip.update("total --") + self._drop_chip.update("dropped --") + self._drop_chip.remove_class("warning-chip") + self._fill_bar.total = 1 + self._fill_bar.progress = 0 + self._fill_text.update("--/--") + return + + stage_1sec = _find_stage_metric(dataloader_1_second, self.stage_key) + stage_total = _find_stage_metric(dataloader_total, self.stage_key) + self._update_name(stage_1sec or stage_total) + + queue_1sec = _get_queue_metric(stage_1sec, self._queue_name) + queue_total = _get_queue_metric(stage_total, self._queue_name) + + queue_name = None + if queue_1sec and queue_1sec.name: + queue_name = queue_1sec.name + elif queue_total and queue_total.name: + queue_name = queue_total.name + elif self._queue_name: + queue_name = self._queue_name + self._queue_name_chip.update( + f"queue {queue_name}" if queue_name else "queue --" + ) + + rate = queue_1sec.get_count if queue_1sec else 0 + self._rate_chip.update(f"rate {format_si(rate)}/s") + if rate == 0: + self._rate_chip.add_class("queue-rate--zero") + else: + self._rate_chip.remove_class("queue-rate--zero") + + if queue_total: + self._total_chip.update( + f"total {format_full_number(queue_total.get_count)}" + ) + else: + self._total_chip.update("total --") + + has_drops = False + if queue_1sec and queue_1sec.drop_count: + self._drop_chip.update( + f"dropped {format_si(queue_1sec.drop_count)}/s" + ) + has_drops = True + elif queue_total and queue_total.drop_count: + self._drop_chip.update( + f"dropped {format_full_number(queue_total.drop_count)}" + ) + has_drops = True + else: + self._drop_chip.update("dropped --") + has_drops = False + + if has_drops: + self._drop_chip.add_class("warning-chip") + else: + self._drop_chip.remove_class("warning-chip") + + size = _average_queue_fullness(queue_1sec) + capacity: int | None = None + if queue_1sec and queue_1sec.queue_capacity > 0: + capacity = queue_1sec.queue_capacity + elif queue_total and queue_total.queue_capacity > 0: + capacity = queue_total.queue_capacity + + if capacity and capacity > 0: + self._fill_bar.total = capacity + if size is not None: + self._fill_bar.progress = min(size, capacity) + fill_text = ( + f"{format_full_number(size)}/{format_full_number(capacity)}" + ) + else: + self._fill_bar.progress = 0 + fill_text = f"--/{format_full_number(capacity)}" + else: + self._fill_bar.total = 1 + self._fill_bar.progress = 0 + fill_text = "--/--" + self._fill_text.update(fill_text) diff --git a/src/lczero_training/tui/log_pane.py b/src/lczero_training/tui/log_pane.py new file mode 100644 index 00000000..b82bb342 --- /dev/null +++ b/src/lczero_training/tui/log_pane.py @@ -0,0 +1,67 @@ +import datetime +from pathlib import Path +from typing import Any, Optional, TextIO + +from anyio.streams.text import TextReceiveStream +from textual.widgets import RichLog + + +class StreamingLogPane(RichLog): + """Log pane that streams output from an async text stream.""" + + def __init__( + self, + stream: TextReceiveStream, + logfile_path: Optional[str] = None, + **kwargs: Any, + ) -> None: + super().__init__( + highlight=True, markup=True, max_lines=1000, wrap=True, **kwargs + ) + self._stream = stream + self._logfile_path = logfile_path + self._logfile: Optional[Path] = None + self._logfile_handle: Optional[TextIO] = None + if logfile_path: + self._logfile = Path(logfile_path) + + def on_mount(self) -> None: + """Start the async reading task when the widget is mounted.""" + if self._logfile: + self._write_banner() + self.run_worker(self._read_stream()) + + def _write_banner(self) -> None: + """Write a session banner to the logfile.""" + if not self._logfile: + return + + self._logfile.parent.mkdir(parents=True, exist_ok=True) + self._logfile_handle = self._logfile.open("a", encoding="utf-8") + timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + banner = ( + f"\n{'=' * 80}\n" + f"LCZero Training TUI Session Started: {timestamp}\n" + f"{'=' * 80}\n" + ) + self._logfile_handle.write(banner) + self._logfile_handle.flush() + + def _write_to_file(self, line: str) -> None: + """Write a line to the logfile.""" + if not self._logfile_handle: + return + + self._logfile_handle.write(f"{line}\n") + self._logfile_handle.flush() + + async def _read_stream(self) -> None: + """Async function that reads lines from the text stream.""" + try: + async for line in self._stream: + line = line.strip() + if line: + self.write(line) + self._write_to_file(line) + except Exception: + pass diff --git a/src/lczero_training/tui/training_widgets.py b/src/lczero_training/tui/training_widgets.py new file mode 100644 index 00000000..b3aeac44 --- /dev/null +++ b/src/lczero_training/tui/training_widgets.py @@ -0,0 +1,118 @@ +from textual.app import ComposeResult +from textual.widgets import ProgressBar, Static + +from ..daemon.protocol.messages import TrainingScheduleData + + +class TimeProgressWidget(Static): + """A widget to display a label, progress bar, and time ratio.""" + + def __init__(self, label: str, *, id: str | None = None) -> None: + super().__init__(id=id) + self._label = label + + def compose(self) -> ComposeResult: + yield Static(self._label, classes="time-label") + yield ProgressBar(show_eta=False, classes="time-progress-bar") + yield Static("", classes="time-ratio") + + def update_progress( + self, + current: float, + total: float, + current_formatted: str | None = None, + total_formatted: str | None = None, + ) -> None: + """Update the progress bar and ratio text.""" + progress_bar = self.query_one(ProgressBar) + + progress_bar.total = total or None + progress_bar.progress = current + + current_str = ( + current_formatted + if current_formatted is not None + else str(int(current)) + ) + total_str = ( + total_formatted if total_formatted is not None else str(int(total)) + ) + + self.query_one(".time-ratio", Static).update( + f"{current_str}/{total_str}" + ) + + +def format_time_duration(seconds: float) -> str: + """Format time duration in seconds to human readable format with days support.""" + if seconds <= 0: + return "--" + + total_seconds = int(seconds) + days, remainder = divmod(total_seconds, 86400) + hours, remainder = divmod(remainder, 3600) + minutes, secs = divmod(remainder, 60) + + if days > 0: + return f"{days}d {hours:02d}:{minutes:02d}:{secs:02d}" + if hours > 0: + return f"{hours:02d}:{minutes:02d}:{secs:02d}" + return f"{minutes:02d}:{secs:02d}" + + +class TrainingScheduleWidget(Static): + """A widget to display training schedule information.""" + + def compose(self) -> ComposeResult: + yield Static("Uptime: -- Stage: --", id="uptime-stage-display") + yield Static("Completed epochs: 0", id="epochs-display") + yield TimeProgressWidget("New Chunks:", id="chunks-progress") + yield TimeProgressWidget("Training time", id="training-time-progress") + yield TimeProgressWidget("Cycle time", id="cycle-time-progress") + + def update_training_schedule( + self, data: TrainingScheduleData | None + ) -> None: + """Update the widget with new training schedule data.""" + if not data: + return + + uptime_str = format_time_duration(data.total_uptime_seconds) + self.query_one("#uptime-stage-display", Static).update( + f"Uptime: {uptime_str} Stage: {data.current_stage.value}" + ) + + self.query_one("#epochs-display", Static).update( + f"Completed epochs: {data.completed_epochs_since_start}" + ) + + self.query_one("#chunks-progress", TimeProgressWidget).update_progress( + current=data.new_chunks_since_training_start, + total=data.chunks_to_wait, + ) + + self.query_one( + "#training-time-progress", TimeProgressWidget + ).update_progress( + current=data.current_training_time_seconds, + total=data.previous_training_time_seconds, + current_formatted=format_time_duration( + data.current_training_time_seconds + ), + total_formatted=format_time_duration( + data.previous_training_time_seconds + ), + ) + + self.query_one( + "#cycle-time-progress", TimeProgressWidget + ).update_progress( + current=data.current_cycle_time_seconds, + total=data.previous_cycle_time_seconds, + current_formatted=format_time_duration( + data.current_cycle_time_seconds + ), + total_formatted=format_time_duration( + data.previous_cycle_time_seconds + ), + ) diff --git a/src/proto/__init__.py b/src/proto/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/subprojects/abseil-cpp.wrap b/subprojects/abseil-cpp.wrap new file mode 100644 index 00000000..9ee5db8e --- /dev/null +++ b/subprojects/abseil-cpp.wrap @@ -0,0 +1,105 @@ +[wrap-file] +directory = abseil-cpp-20250814.1 +source_url = https://github.com/abseil/abseil-cpp/releases/download/20250814.1/abseil-cpp-20250814.1.tar.gz +source_filename = abseil-cpp-20250814.1.tar.gz +source_hash = 1692f77d1739bacf3f94337188b78583cf09bab7e420d2dc6c5605a4f86785a1 +patch_directory = abseil-cpp-20250814.1 +source_fallback_url = https://github.com/mesonbuild/wrapdb/releases/download/abseil-cpp_20250814.1-1/abseil-cpp-20250814.1.tar.gz + +[provide] +absl_base = absl_base_dep +absl_container = absl_container_dep +absl_debugging = absl_debugging_dep +absl_log = absl_log_dep +absl_flags = absl_flags_dep +absl_hash = absl_hash_dep +absl_crc = absl_crc_dep +absl_numeric = absl_numeric_dep +absl_profiling = absl_profiling_dep +absl_random = absl_random_dep +absl_status = absl_status_dep +absl_strings = absl_strings_dep +absl_synchronization = absl_synchronization_dep +absl_time = absl_time_dep +absl_types = absl_types_dep +absl_algorithm_container = absl_base_dep +absl_any_invocable = absl_base_dep +absl_bad_any_cast_impl = absl_types_dep +absl_bad_optional_access = absl_types_dep +absl_bad_variant_access = absl_types_dep +absl_bind_front = absl_base_dep +absl_city = absl_hash_dep +absl_civil_time = absl_time_dep +absl_cleanup = absl_base_dep +absl_cord = absl_strings_dep +absl_cord_internal = absl_strings_dep +absl_cordz_functions = absl_strings_dep +absl_cordz_handle = absl_strings_dep +absl_cordz_info = absl_strings_dep +absl_cordz_sample_token = absl_strings_dep +absl_core_headers = absl_base_dep +absl_crc32c = absl_crc_dep +absl_debugging_internal = absl_debugging_dep +absl_demangle_internal = absl_debugging_dep +absl_die_if_null = absl_log_dep +absl_examine_stack = absl_debugging_dep +absl_exponential_biased = absl_profiling_dep +absl_failure_signal_handler = absl_debugging_dep +absl_flags_commandlineflag = absl_flags_dep +absl_flags_commandlineflag_internal = absl_flags_dep +absl_flags_config = absl_flags_dep +absl_flags_internal = absl_flags_dep +absl_flags_marshalling = absl_flags_dep +absl_flags_parse = absl_flags_dep +absl_flags_private_handle_accessor = absl_flags_dep +absl_flags_program_name = absl_flags_dep +absl_flags_reflection = absl_flags_dep +absl_flags_usage = absl_flags_dep +absl_flags_usage_internal = absl_flags_dep +absl_flat_hash_map = absl_container_dep +absl_flat_hash_set = absl_container_dep +absl_function_ref = absl_base_dep +absl_graphcycles_internal = absl_synchronization_dep +absl_hashtablez_sampler = absl_container_dep +absl_inlined_vector = absl_container_dep +absl_int128 = absl_numeric_dep +absl_leak_check = absl_debugging_dep +absl_log_initialize = absl_log_dep +absl_log_internal_check_op = absl_log_dep +absl_log_internal_message = absl_log_dep +absl_log_severity = absl_base_dep +absl_low_level_hash = absl_hash_dep +absl_memory = absl_base_dep +absl_optional = absl_types_dep +absl_periodic_sampler = absl_profiling_dep +absl_random_bit_gen_ref = absl_random_dep +absl_random_distributions = absl_random_dep +absl_random_internal_distribution_test_util = absl_random_dep +absl_random_internal_platform = absl_random_dep +absl_random_internal_pool_urbg = absl_random_dep +absl_random_internal_randen = absl_random_dep +absl_random_internal_randen_hwaes = absl_random_dep +absl_random_internal_randen_hwaes_impl = absl_random_dep +absl_random_internal_randen_slow = absl_random_dep +absl_random_internal_seed_material = absl_random_dep +absl_random_random = absl_random_dep +absl_random_seed_gen_exception = absl_random_dep +absl_random_seed_sequences = absl_random_dep +absl_raw_hash_set = absl_container_dep +absl_raw_logging_internal = absl_base_dep +absl_scoped_set_env = absl_base_dep +absl_span = absl_types_dep +absl_spinlock_wait = absl_base_dep +absl_stacktrace = absl_debugging_dep +absl_statusor = absl_status_dep +absl_str_format = absl_strings_dep +absl_str_format_internal = absl_strings_dep +absl_strerror = absl_base_dep +absl_string_view = absl_strings_dep +absl_strings_internal = absl_strings_dep +absl_symbolize = absl_debugging_dep +absl_throw_delegate = absl_base_dep +absl_time_zone = absl_time_dep +absl_type_traits = absl_base_dep +absl_utility = absl_base_dep +absl_variant = absl_types_dep diff --git a/subprojects/gaviotatb.wrap b/subprojects/gaviotatb.wrap new file mode 100644 index 00000000..8bd32f9d --- /dev/null +++ b/subprojects/gaviotatb.wrap @@ -0,0 +1,7 @@ +[wrap-git] +url = https://github.com/michiguel/Gaviota-Tablebases +revision = head + +patch_directory = gaviotatb + + diff --git a/subprojects/gtest.wrap b/subprojects/gtest.wrap new file mode 100644 index 00000000..9902a4f7 --- /dev/null +++ b/subprojects/gtest.wrap @@ -0,0 +1,16 @@ +[wrap-file] +directory = googletest-1.17.0 +source_url = https://github.com/google/googletest/archive/refs/tags/v1.17.0.tar.gz +source_filename = googletest-1.17.0.tar.gz +source_hash = 65fab701d9829d38cb77c14acdc431d2108bfdbf8979e40eb8ae567edf10b27c +patch_filename = gtest_1.17.0-4_patch.zip +patch_url = https://wrapdb.mesonbuild.com/v2/gtest_1.17.0-4/get_patch +patch_hash = 3abf7662d09db706453a5b064a1e914678c74b9d9b0b19382747ca561d0d8750 +source_fallback_url = https://github.com/mesonbuild/wrapdb/releases/download/gtest_1.17.0-4/googletest-1.17.0.tar.gz +wrapdb_version = 1.17.0-4 + +[provide] +gtest = gtest_dep +gtest_main = gtest_main_dep +gmock = gmock_dep +gmock_main = gmock_main_dep diff --git a/subprojects/packagefiles/abseil-cpp-20250814.1/LICENSE.build b/subprojects/packagefiles/abseil-cpp-20250814.1/LICENSE.build new file mode 100644 index 00000000..b59833de --- /dev/null +++ b/subprojects/packagefiles/abseil-cpp-20250814.1/LICENSE.build @@ -0,0 +1,19 @@ +Copyright (c) 2021 The Meson development team + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/subprojects/packagefiles/abseil-cpp-20250814.1/meson.build b/subprojects/packagefiles/abseil-cpp-20250814.1/meson.build new file mode 100644 index 00000000..33fbc813 --- /dev/null +++ b/subprojects/packagefiles/abseil-cpp-20250814.1/meson.build @@ -0,0 +1,887 @@ +project( + 'abseil-cpp', + 'cpp', + meson_version: '>=0.49.0', + version: '20250814.1', + license: 'Apache-2.0', +) + +override_cpp = 'cpp_std=c++17' + +cpp = meson.get_compiler('cpp') + +flags = cpp.get_supported_arguments( + '/DNOMINMAX', + '-Wcast-qual', + '-Wconversion-null', + '-Wmissing-declarations', + '-Wno-comma', + '-Wno-conversion', + '-Wno-double-promotion', + '-Wno-float-promotion', + '-Wno-format-literal', + '-Wno-gcc-compat', + '-Wno-old-style-cast', + '-Wno-packed', + '-Wno-padded', + '-Wno-pedantic', + '-Wno-range-loop-analysis', + '-Wno-sign-compare', + '-Woverlength-strings', + '-Wpointer-arith', + '-Wswitch-enums', + '-Wunused-local-typedefs', + '-Wunused-result', + '-Wvarargs', + '-Wvla', + '-Wwrite-strings', +) + +add_project_arguments( + flags, + language: 'cpp', +) + +arch_cpp_flags = [] +hw_cpp_flags = [] +unscaled_cycleclock_flag = [] +if host_machine.cpu_family() == 'x86_64' + hw_cpp_flags += ['-maes', '-msse4.1'] +elif host_machine.cpu_family() == 'aarch64' and cpp.sizeof('void*') == 8 + hw_cpp_flags += ['-march=armv8-a+crypto'] +elif host_machine.cpu_family() == 'arm' and cpp.sizeof('void*') == 4 + hw_cpp_flags += ['-mfpu=neon'] +elif host_machine.cpu_family() == 'ppc' or host_machine.cpu_family() == 'ppc64' + # This will work with glibc but not musl + timebase_check = '''#include + int main() { + __ppc_get_timebase_freq(); + return 0; + }''' + if not cpp.compiles(timebase_check) + unscaled_cycleclock_flag += ['-DABSL_USE_UNSCALED_CYCLECLOCK=0'] + endif +endif +arch_flags = cpp.get_supported_arguments(arch_cpp_flags) + unscaled_cycleclock_flag +hw_flags = cpp.get_supported_arguments(hw_cpp_flags) + +libatomic = dependency( + '', + required: false, +) +if cpp.get_argument_syntax() != 'msvc' and not cpp.links( + 'int main(){__sync_synchronize();}', + name: 'atomic builtins', +) + libatomic = cpp.find_library('atomic') +endif + +absl_include_dir = include_directories('.') + +# Group files by the containing library +absl_base_sources = files( + 'absl/base/internal/cycleclock.cc', + 'absl/base/internal/low_level_alloc.cc', + 'absl/base/internal/poison.cc', + 'absl/base/internal/raw_logging.cc', + 'absl/base/internal/scoped_set_env.cc', + 'absl/base/internal/spinlock.cc', + 'absl/base/internal/spinlock_wait.cc', + 'absl/base/internal/strerror.cc', + 'absl/base/internal/sysinfo.cc', + 'absl/base/internal/thread_identity.cc', + 'absl/base/internal/throw_delegate.cc', + 'absl/base/internal/tracing.cc', + 'absl/base/internal/unscaledcycleclock.cc', + 'absl/base/log_severity.cc', +) +absl_base_headers = files( + 'absl/base/attributes.h', + 'absl/base/call_once.h', + 'absl/base/casts.h', + 'absl/base/config.h', + 'absl/base/const_init.h', + 'absl/base/dynamic_annotations.h', + 'absl/base/fast_type_id.h', + 'absl/base/internal/atomic_hook.h', + 'absl/base/internal/atomic_hook_test_helper.h', + 'absl/base/internal/cycleclock.h', + 'absl/base/internal/cycleclock_config.h', + 'absl/base/internal/direct_mmap.h', + 'absl/base/internal/dynamic_annotations.h', + 'absl/base/internal/endian.h', + 'absl/base/internal/errno_saver.h', + 'absl/base/internal/exception_safety_testing.h', + 'absl/base/internal/exception_testing.h', + 'absl/base/internal/hide_ptr.h', + 'absl/base/internal/identity.h', + 'absl/base/internal/low_level_alloc.h', + 'absl/base/internal/low_level_scheduling.h', + 'absl/base/internal/per_thread_tls.h', + 'absl/base/internal/poison.h', + 'absl/base/internal/pretty_function.h', + 'absl/base/internal/raw_logging.h', + 'absl/base/internal/scheduling_mode.h', + 'absl/base/internal/scoped_set_env.h', + 'absl/base/internal/spinlock.h', + 'absl/base/internal/spinlock_akaros.inc', + 'absl/base/internal/spinlock_linux.inc', + 'absl/base/internal/spinlock_posix.inc', + 'absl/base/internal/spinlock_wait.h', + 'absl/base/internal/spinlock_win32.inc', + 'absl/base/internal/strerror.h', + 'absl/base/internal/sysinfo.h', + 'absl/base/internal/thread_identity.h', + 'absl/base/internal/throw_delegate.h', + 'absl/base/internal/tsan_mutex_interface.h', + 'absl/base/internal/tracing.h', + 'absl/base/internal/unaligned_access.h', + 'absl/base/internal/unscaledcycleclock.h', + 'absl/base/internal/unscaledcycleclock_config.h', + 'absl/base/log_severity.h', + 'absl/base/macros.h', + 'absl/base/no_destructor.h', + 'absl/base/nullability.h', + 'absl/base/optimization.h', + 'absl/base/options.h', + 'absl/base/policy_checks.h', + 'absl/base/port.h', + 'absl/base/thread_annotations.h', + 'absl/functional/any_invocable.h', + 'absl/functional/internal/any_invocable.h', + 'absl/memory/memory.h', + 'absl/meta/type_traits.h', + 'absl/utility/utility.h', + # Dependent headers of absl_base +) + +absl_container_sources = files( + 'absl/container/internal/hashtablez_sampler.cc', + 'absl/container/internal/hashtablez_sampler_force_weak_definition.cc', + 'absl/container/internal/raw_hash_set.cc', +) +absl_container_headers = files( + 'absl/container/btree_map.h', + 'absl/container/btree_set.h', + 'absl/container/btree_test.h', + 'absl/container/fixed_array.h', + 'absl/container/flat_hash_map.h', + 'absl/container/flat_hash_set.h', + 'absl/container/hash_container_defaults.h', + 'absl/container/inlined_vector.h', + 'absl/container/internal/btree.h', + 'absl/container/internal/btree_container.h', + 'absl/container/internal/common.h', + 'absl/container/internal/common_policy_traits.h', + 'absl/container/internal/compressed_tuple.h', + 'absl/container/internal/container_memory.h', + 'absl/container/internal/hash_function_defaults.h', + 'absl/container/internal/hash_generator_testing.h', + 'absl/container/internal/hash_policy_testing.h', + 'absl/container/internal/hash_policy_traits.h', + 'absl/container/internal/hashtable_debug.h', + 'absl/container/internal/hashtable_debug_hooks.h', + 'absl/container/internal/hashtablez_sampler.h', + 'absl/container/internal/inlined_vector.h', + 'absl/container/internal/layout.h', + 'absl/container/internal/node_slot_policy.h', + 'absl/container/internal/raw_hash_map.h', + 'absl/container/internal/raw_hash_set.h', + 'absl/container/internal/test_instance_tracker.h', + 'absl/container/internal/tracked.h', + 'absl/container/internal/unordered_map_constructor_test.h', + 'absl/container/internal/unordered_map_lookup_test.h', + 'absl/container/internal/unordered_map_members_test.h', + 'absl/container/internal/unordered_map_modifiers_test.h', + 'absl/container/internal/unordered_set_constructor_test.h', + 'absl/container/internal/unordered_set_lookup_test.h', + 'absl/container/internal/unordered_set_members_test.h', + 'absl/container/internal/unordered_set_modifiers_test.h', + 'absl/container/node_hash_map.h', + 'absl/container/node_hash_set.h', +) + +absl_crc_sources = files( + 'absl/crc/crc32c.cc', + 'absl/crc/internal/cpu_detect.cc', + 'absl/crc/internal/crc.cc', + 'absl/crc/internal/crc_cord_state.cc', + 'absl/crc/internal/crc_memcpy_fallback.cc', + 'absl/crc/internal/crc_memcpy_x86_arm_combined.cc', + 'absl/crc/internal/crc_non_temporal_memcpy.cc', + 'absl/crc/internal/crc_x86_arm_combined.cc', +) +absl_crc_headers = files( + 'absl/crc/crc32c.h', + 'absl/crc/internal/cpu_detect.h', + 'absl/crc/internal/crc.h', + 'absl/crc/internal/crc32_x86_arm_combined_simd.h', + 'absl/crc/internal/crc32c.h', + 'absl/crc/internal/crc32c_inline.h', + 'absl/crc/internal/crc_cord_state.h', + 'absl/crc/internal/crc_internal.h', + 'absl/crc/internal/crc_memcpy.h', + 'absl/crc/internal/non_temporal_arm_intrinsics.h', + 'absl/crc/internal/non_temporal_memcpy.h', +) + +absl_debugging_sources = files( + 'absl/debugging/failure_signal_handler.cc', + 'absl/debugging/internal/address_is_readable.cc', + 'absl/debugging/internal/decode_rust_punycode.cc', + 'absl/debugging/internal/demangle.cc', + 'absl/debugging/internal/demangle_rust.cc', + 'absl/debugging/internal/elf_mem_image.cc', + 'absl/debugging/internal/examine_stack.cc', + 'absl/debugging/internal/stack_consumption.cc', + 'absl/debugging/internal/utf8_for_code_point.cc', + 'absl/debugging/internal/vdso_support.cc', + 'absl/debugging/leak_check.cc', + 'absl/debugging/stacktrace.cc', + 'absl/debugging/symbolize.cc', +) +absl_debugging_headers = files( + 'absl/debugging/failure_signal_handler.h', + 'absl/debugging/internal/address_is_readable.h', + 'absl/debugging/internal/bounded_utf8_length_sequence.h', + 'absl/debugging/internal/decode_rust_punycode.h', + 'absl/debugging/internal/demangle.h', + 'absl/debugging/internal/demangle_rust.h', + 'absl/debugging/internal/elf_mem_image.h', + 'absl/debugging/internal/examine_stack.h', + 'absl/debugging/internal/stack_consumption.h', + 'absl/debugging/internal/stacktrace_aarch64-inl.inc', + 'absl/debugging/internal/stacktrace_arm-inl.inc', + 'absl/debugging/internal/stacktrace_config.h', + 'absl/debugging/internal/stacktrace_emscripten-inl.inc', + 'absl/debugging/internal/stacktrace_generic-inl.inc', + 'absl/debugging/internal/stacktrace_powerpc-inl.inc', + 'absl/debugging/internal/stacktrace_riscv-inl.inc', + 'absl/debugging/internal/stacktrace_unimplemented-inl.inc', + 'absl/debugging/internal/stacktrace_win32-inl.inc', + 'absl/debugging/internal/stacktrace_x86-inl.inc', + 'absl/debugging/internal/symbolize.h', + 'absl/debugging/internal/utf8_for_code_point.h', + 'absl/debugging/internal/vdso_support.h', + 'absl/debugging/leak_check.h', + 'absl/debugging/stacktrace.h', + 'absl/debugging/symbolize.h', + 'absl/debugging/symbolize_darwin.inc', + 'absl/debugging/symbolize_elf.inc', + 'absl/debugging/symbolize_emscripten.inc', + 'absl/debugging/symbolize_unimplemented.inc', + 'absl/debugging/symbolize_win32.inc', +) + +absl_flags_sources = files( + 'absl/flags/commandlineflag.cc', + 'absl/flags/internal/commandlineflag.cc', + 'absl/flags/internal/flag.cc', + 'absl/flags/internal/flag.cc', + 'absl/flags/internal/private_handle_accessor.cc', + 'absl/flags/internal/program_name.cc', + 'absl/flags/internal/usage.cc', + 'absl/flags/marshalling.cc', + 'absl/flags/parse.cc', + 'absl/flags/reflection.cc', + 'absl/flags/usage.cc', + 'absl/flags/usage_config.cc', +) +absl_flags_headers = files( + 'absl/flags/commandlineflag.h', + 'absl/flags/config.h', + 'absl/flags/declare.h', + 'absl/flags/internal/commandlineflag.h', + 'absl/flags/internal/flag.h', + 'absl/flags/internal/flag.h', + 'absl/flags/internal/parse.h', + 'absl/flags/internal/path_util.h', + 'absl/flags/internal/private_handle_accessor.h', + 'absl/flags/internal/program_name.h', + 'absl/flags/internal/registry.h', + 'absl/flags/internal/sequence_lock.h', + 'absl/flags/internal/usage.h', + 'absl/flags/marshalling.h', + 'absl/flags/parse.h', + 'absl/flags/reflection.h', + 'absl/flags/usage.h', + 'absl/flags/usage_config.h', +) + +absl_hash_sources = files( + 'absl/hash/internal/city.cc', + 'absl/hash/internal/hash.cc', +) +absl_hash_headers = files( + 'absl/hash/hash.h', + 'absl/hash/hash_testing.h', + 'absl/hash/internal/city.h', + 'absl/hash/internal/hash.h', + 'absl/hash/internal/spy_hash_state.h', +) + +absl_log_sources = files( + 'absl/log/die_if_null.cc', + 'absl/log/flags.cc', + 'absl/log/globals.cc', + 'absl/log/initialize.cc', + 'absl/log/internal/check_op.cc', + 'absl/log/internal/conditions.cc', + 'absl/log/internal/fnmatch.cc', + 'absl/log/internal/globals.cc', + 'absl/log/internal/log_format.cc', + 'absl/log/internal/log_message.cc', + 'absl/log/internal/log_sink_set.cc', + 'absl/log/internal/nullguard.cc', + 'absl/log/internal/proto.cc', + 'absl/log/internal/structured_proto.cc', + 'absl/log/internal/vlog_config.cc', + 'absl/log/log_entry.cc', + 'absl/log/log_sink.cc', +) +absl_log_headers = files( + 'absl/log/absl_check.h', + 'absl/log/absl_log.h', + 'absl/log/check.h', + 'absl/log/die_if_null.h', + 'absl/log/flags.h', + 'absl/log/globals.h', + 'absl/log/initialize.h', + 'absl/log/internal/append_truncated.h', + 'absl/log/internal/check_impl.h', + 'absl/log/internal/check_op.h', + 'absl/log/internal/conditions.h', + 'absl/log/internal/config.h', + 'absl/log/internal/flags.h', + 'absl/log/internal/fnmatch.h', + 'absl/log/internal/globals.h', + 'absl/log/internal/log_format.h', + 'absl/log/internal/log_impl.h', + 'absl/log/internal/log_message.h', + 'absl/log/internal/log_sink_set.h', + 'absl/log/internal/nullguard.h', + 'absl/log/internal/nullstream.h', + 'absl/log/internal/proto.h', + 'absl/log/internal/strip.h', + 'absl/log/internal/structured.h', + 'absl/log/internal/structured_proto.h', + 'absl/log/internal/test_actions.h', + 'absl/log/internal/test_helpers.h', + 'absl/log/internal/test_matchers.h', + 'absl/log/internal/vlog_config.h', + 'absl/log/internal/voidify.h', + 'absl/log/log.h', + 'absl/log/log_entry.h', + 'absl/log/log_sink.h', + 'absl/log/log_sink_registry.h', + 'absl/log/log_streamer.h', + 'absl/log/scoped_mock_log.h', + 'absl/log/structured.h', +) + +absl_numeric_sources = files('absl/numeric/int128.cc') +absl_numeric_headers = files( + 'absl/numeric/bits.h', + 'absl/numeric/int128.h', + 'absl/numeric/int128_have_intrinsic.inc', + 'absl/numeric/int128_no_intrinsic.inc', + 'absl/numeric/internal/bits.h', + 'absl/numeric/internal/representation.h', +) + +absl_profiling_sources = files( + 'absl/profiling/internal/exponential_biased.cc', + 'absl/profiling/internal/periodic_sampler.cc', +) +absl_profiling_headers = files( + 'absl/profiling/internal/exponential_biased.h', + 'absl/profiling/internal/periodic_sampler.h', + 'absl/profiling/internal/sample_recorder.h', +) + +absl_random_sources = files( + 'absl/random/discrete_distribution.cc', + 'absl/random/gaussian_distribution.cc', + 'absl/random/internal/chi_square.cc', + 'absl/random/internal/entropy_pool.cc', + 'absl/random/internal/randen.cc', + 'absl/random/internal/randen_detect.cc', + 'absl/random/internal/randen_hwaes.cc', + 'absl/random/internal/randen_round_keys.cc', + 'absl/random/internal/randen_slow.cc', + 'absl/random/internal/seed_material.cc', + 'absl/random/seed_gen_exception.cc', + 'absl/random/seed_sequences.cc', +) +absl_random_headers = files( + 'absl/random/bernoulli_distribution.h', + 'absl/random/beta_distribution.h', + 'absl/random/bit_gen_ref.h', + 'absl/random/discrete_distribution.h', + 'absl/random/distributions.h', + 'absl/random/exponential_distribution.h', + 'absl/random/gaussian_distribution.h', + 'absl/random/internal/chi_square.h', + 'absl/random/internal/distribution_caller.h', + 'absl/random/internal/distribution_test_util.h', + 'absl/random/internal/entropy_pool.h', + 'absl/random/internal/explicit_seed_seq.h', + 'absl/random/internal/fast_uniform_bits.h', + 'absl/random/internal/fastmath.h', + 'absl/random/internal/generate_real.h', + 'absl/random/internal/iostream_state_saver.h', + 'absl/random/internal/mock_helpers.h', + 'absl/random/internal/mock_overload_set.h', + 'absl/random/internal/nanobenchmark.h', + 'absl/random/internal/nonsecure_base.h', + 'absl/random/internal/pcg_engine.h', + 'absl/random/internal/platform.h', + 'absl/random/internal/randen.h', + 'absl/random/internal/randen_detect.h', + 'absl/random/internal/randen_engine.h', + 'absl/random/internal/randen_hwaes.h', + 'absl/random/internal/randen_slow.h', + 'absl/random/internal/randen_traits.h', + 'absl/random/internal/salted_seed_seq.h', + 'absl/random/internal/seed_material.h', + 'absl/random/internal/sequence_urbg.h', + 'absl/random/internal/traits.h', + 'absl/random/internal/uniform_helper.h', + 'absl/random/internal/wide_multiply.h', + 'absl/random/log_uniform_int_distribution.h', + 'absl/random/mock_distributions.h', + 'absl/random/mocking_bit_gen.h', + 'absl/random/poisson_distribution.h', + 'absl/random/random.h', + 'absl/random/seed_gen_exception.h', + 'absl/random/seed_sequences.h', + 'absl/random/uniform_int_distribution.h', + 'absl/random/uniform_real_distribution.h', + 'absl/random/zipf_distribution.h', +) + +absl_status_sources = files( + 'absl/status/internal/status_internal.cc', + 'absl/status/status.cc', + 'absl/status/status_payload_printer.cc', + 'absl/status/statusor.cc', +) +absl_status_headers = files( + 'absl/status/internal/status_internal.h', + 'absl/status/internal/statusor_internal.h', + 'absl/status/status.h', + 'absl/status/status_payload_printer.h', + 'absl/status/statusor.h', +) + +absl_strings_sources = files( + 'absl/strings/ascii.cc', + 'absl/strings/charconv.cc', + 'absl/strings/cord.cc', + 'absl/strings/cord_analysis.cc', + 'absl/strings/escaping.cc', + 'absl/strings/internal/charconv_bigint.cc', + 'absl/strings/internal/charconv_parse.cc', + 'absl/strings/internal/cord_internal.cc', + 'absl/strings/internal/cord_rep_btree.cc', + 'absl/strings/internal/cord_rep_btree_navigator.cc', + 'absl/strings/internal/cord_rep_btree_reader.cc', + 'absl/strings/internal/cord_rep_consume.cc', + 'absl/strings/internal/cord_rep_crc.cc', + 'absl/strings/internal/cordz_functions.cc', + 'absl/strings/internal/cordz_handle.cc', + 'absl/strings/internal/cordz_info.cc', + 'absl/strings/internal/cordz_sample_token.cc', + 'absl/strings/internal/damerau_levenshtein_distance.cc', + 'absl/strings/internal/escaping.cc', + 'absl/strings/internal/memutil.cc', + 'absl/strings/internal/ostringstream.cc', + 'absl/strings/internal/pow10_helper.cc', + 'absl/strings/internal/str_format/arg.cc', + 'absl/strings/internal/str_format/bind.cc', + 'absl/strings/internal/str_format/extension.cc', + 'absl/strings/internal/str_format/float_conversion.cc', + 'absl/strings/internal/str_format/output.cc', + 'absl/strings/internal/str_format/parser.cc', + 'absl/strings/internal/stringify_sink.cc', + 'absl/strings/internal/utf8.cc', + 'absl/strings/match.cc', + 'absl/strings/numbers.cc', + 'absl/strings/str_cat.cc', + 'absl/strings/str_replace.cc', + 'absl/strings/str_split.cc', + 'absl/strings/string_view.cc', + 'absl/strings/substitute.cc', +) +absl_strings_headers = files( + 'absl/strings/ascii.h', + 'absl/strings/charconv.h', + 'absl/strings/cord.h', + 'absl/strings/cord_analysis.h', + 'absl/strings/cord_buffer.h', + 'absl/strings/cord_test_helpers.h', + 'absl/strings/cordz_test_helpers.h', + 'absl/strings/escaping.h', + 'absl/strings/has_absl_stringify.h', + 'absl/strings/internal/charconv_bigint.h', + 'absl/strings/internal/charconv_parse.h', + 'absl/strings/internal/cord_data_edge.h', + 'absl/strings/internal/cord_internal.h', + 'absl/strings/internal/cord_rep_btree.h', + 'absl/strings/internal/cord_rep_btree_navigator.h', + 'absl/strings/internal/cord_rep_btree_reader.h', + 'absl/strings/internal/cord_rep_consume.h', + 'absl/strings/internal/cord_rep_crc.h', + 'absl/strings/internal/cord_rep_flat.h', + 'absl/strings/internal/cord_rep_test_util.h', + 'absl/strings/internal/cordz_functions.h', + 'absl/strings/internal/cordz_handle.h', + 'absl/strings/internal/cordz_info.h', + 'absl/strings/internal/cordz_sample_token.h', + 'absl/strings/internal/cordz_statistics.h', + 'absl/strings/internal/cordz_update_scope.h', + 'absl/strings/internal/cordz_update_tracker.h', + 'absl/strings/internal/damerau_levenshtein_distance.h', + 'absl/strings/internal/escaping.h', + 'absl/strings/internal/escaping_test_common.h', + 'absl/strings/internal/memutil.h', + 'absl/strings/internal/numbers_test_common.h', + 'absl/strings/internal/ostringstream.h', + 'absl/strings/internal/pow10_helper.h', + 'absl/strings/internal/resize_uninitialized.h', + 'absl/strings/internal/stl_type_traits.h', + 'absl/strings/internal/str_format/arg.h', + 'absl/strings/internal/str_format/bind.h', + 'absl/strings/internal/str_format/checker.h', + 'absl/strings/internal/str_format/constexpr_parser.h', + 'absl/strings/internal/str_format/extension.h', + 'absl/strings/internal/str_format/float_conversion.h', + 'absl/strings/internal/str_format/output.h', + 'absl/strings/internal/str_format/parser.h', + 'absl/strings/internal/str_join_internal.h', + 'absl/strings/internal/str_split_internal.h', + 'absl/strings/internal/string_constant.h', + 'absl/strings/internal/stringify_sink.h', + 'absl/strings/internal/utf8.h', + 'absl/strings/match.h', + 'absl/strings/numbers.h', + 'absl/strings/str_cat.h', + 'absl/strings/str_format.h', + 'absl/strings/str_join.h', + 'absl/strings/str_replace.h', + 'absl/strings/str_split.h', + 'absl/strings/string_view.h', + 'absl/strings/strip.h', + 'absl/strings/substitute.h', +) + +absl_synchronization_sources = files( + 'absl/synchronization/barrier.cc', + 'absl/synchronization/blocking_counter.cc', + 'absl/synchronization/internal/create_thread_identity.cc', + 'absl/synchronization/internal/futex_waiter.cc', + 'absl/synchronization/internal/graphcycles.cc', + 'absl/synchronization/internal/kernel_timeout.cc', + 'absl/synchronization/internal/per_thread_sem.cc', + 'absl/synchronization/internal/pthread_waiter.cc', + 'absl/synchronization/internal/sem_waiter.cc', + 'absl/synchronization/internal/stdcpp_waiter.cc', + 'absl/synchronization/internal/waiter_base.cc', + 'absl/synchronization/internal/win32_waiter.cc', + 'absl/synchronization/mutex.cc', + 'absl/synchronization/notification.cc', +) +absl_synchronization_headers = files( + 'absl/synchronization/barrier.h', + 'absl/synchronization/blocking_counter.h', + 'absl/synchronization/internal/create_thread_identity.h', + 'absl/synchronization/internal/futex.h', + 'absl/synchronization/internal/graphcycles.h', + 'absl/synchronization/internal/kernel_timeout.h', + 'absl/synchronization/internal/per_thread_sem.h', + 'absl/synchronization/internal/thread_pool.h', + 'absl/synchronization/internal/waiter.h', + 'absl/synchronization/mutex.h', + 'absl/synchronization/notification.h', +) + +absl_time_sources = files( + 'absl/time/civil_time.cc', + 'absl/time/clock.cc', + 'absl/time/duration.cc', + 'absl/time/format.cc', + 'absl/time/internal/cctz/src/civil_time_detail.cc', + 'absl/time/internal/cctz/src/time_zone_fixed.cc', + 'absl/time/internal/cctz/src/time_zone_format.cc', + 'absl/time/internal/cctz/src/time_zone_if.cc', + 'absl/time/internal/cctz/src/time_zone_impl.cc', + 'absl/time/internal/cctz/src/time_zone_info.cc', + 'absl/time/internal/cctz/src/time_zone_libc.cc', + 'absl/time/internal/cctz/src/time_zone_lookup.cc', + 'absl/time/internal/cctz/src/time_zone_posix.cc', + 'absl/time/internal/cctz/src/zone_info_source.cc', + 'absl/time/time.cc', +) +absl_time_headers = files( + 'absl/time/civil_time.h', + 'absl/time/clock.h', + 'absl/time/internal/cctz/include/cctz/civil_time.h', + 'absl/time/internal/cctz/include/cctz/civil_time_detail.h', + 'absl/time/internal/cctz/include/cctz/time_zone.h', + 'absl/time/internal/cctz/include/cctz/zone_info_source.h', + 'absl/time/internal/cctz/src/time_zone_fixed.h', + 'absl/time/internal/cctz/src/time_zone_if.h', + 'absl/time/internal/cctz/src/time_zone_impl.h', + 'absl/time/internal/cctz/src/time_zone_info.h', + 'absl/time/internal/cctz/src/time_zone_libc.h', + 'absl/time/internal/cctz/src/time_zone_posix.h', + 'absl/time/internal/cctz/src/tzfile.h', + 'absl/time/internal/get_current_time_chrono.inc', + 'absl/time/internal/get_current_time_posix.inc', + 'absl/time/internal/test_util.h', + 'absl/time/time.h', +) + +absl_types_sources = files() +absl_types_headers = files( + 'absl/types/any.h', + 'absl/types/compare.h', + 'absl/types/internal/span.h', + 'absl/types/optional.h', + 'absl/types/span.h', + 'absl/types/variant.h', +) + +# Libraries +absl_base_lib = static_library( + 'absl_base', + absl_base_sources, + include_directories: absl_include_dir, + override_options: override_cpp, + cpp_args: arch_flags, + dependencies: [dependency('threads'), libatomic], +) + +absl_hash_lib = static_library( + 'absl_hash', + absl_hash_sources, + include_directories: absl_include_dir, + override_options: override_cpp, +) + +absl_numeric_lib = static_library( + 'absl_numeric', + absl_numeric_sources, + include_directories: absl_include_dir, + override_options: override_cpp, +) + +absl_profiling_lib = static_library( + 'absl_profiling', + absl_profiling_sources, + include_directories: absl_include_dir, + override_options: override_cpp, +) + +absl_crc_lib = static_library( + 'absl_crc', + absl_crc_sources, + include_directories: absl_include_dir, + override_options: override_cpp, + link_with: [absl_base_lib], + dependencies: libatomic, +) + +absl_strings_lib = static_library( + 'absl_strings', + absl_strings_sources, + include_directories: absl_include_dir, + override_options: override_cpp, + link_with: [absl_base_lib, absl_crc_lib, absl_numeric_lib, absl_profiling_lib], +) + +absl_debugging_lib = static_library( + 'absl_debugging', + absl_debugging_sources, + include_directories: absl_include_dir, + override_options: override_cpp, + link_with: [absl_base_lib, absl_strings_lib], + dependencies: libatomic, +) + +absl_random_lib = static_library( + 'absl_random', + absl_random_sources, + include_directories: absl_include_dir, + override_options: override_cpp, + cpp_args: hw_flags, + link_with: [absl_base_lib, absl_strings_lib], + dependencies: libatomic, +) + +absl_time_lib = static_library( + 'absl_time', + absl_time_sources, + include_directories: absl_include_dir, + override_options: override_cpp, + link_with: [absl_base_lib, absl_numeric_lib, absl_strings_lib], + # macOS only, upstream: https://github.com/abseil/abseil-cpp/pull/280 + dependencies: dependency( + 'appleframeworks', + modules: 'CoreFoundation', + required: host_machine.system() == 'darwin', + ), +) + +absl_synchronization_lib = static_library( + 'absl_synchronization', + absl_synchronization_sources, + cpp_args: unscaled_cycleclock_flag, + include_directories: absl_include_dir, + override_options: override_cpp, + link_with: [absl_base_lib, absl_debugging_lib, absl_time_lib], +) + +absl_container_lib = static_library( + 'absl_container', + absl_container_sources, + include_directories: absl_include_dir, + override_options: override_cpp, + link_with: [ + absl_base_lib, + absl_debugging_lib, + absl_hash_lib, + absl_synchronization_lib, + absl_time_lib, + ], +) + +absl_flags_lib = static_library( + 'absl_flags', + absl_flags_sources, + include_directories: absl_include_dir, + override_options: override_cpp, + link_with: [ + absl_base_lib, + absl_container_lib, + absl_hash_lib, + absl_strings_lib, + absl_synchronization_lib, + ], + dependencies: libatomic, +) + +absl_status_lib = static_library( + 'absl_status', + absl_status_sources, + include_directories: absl_include_dir, + override_options: override_cpp, + link_with: [absl_base_lib, absl_strings_lib], +) + +absl_log_lib = static_library( + 'absl_log', + absl_log_sources, + cpp_args: unscaled_cycleclock_flag, + include_directories: absl_include_dir, + override_options: override_cpp, + link_with: [absl_base_lib, absl_strings_lib, absl_flags_lib], + dependencies: libatomic, +) + +# Dependencies +absl_base_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_base_lib, +) + +absl_hash_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_hash_lib, +) + +absl_numeric_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_numeric_lib, +) + +absl_profiling_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_profiling_lib, +) + +absl_strings_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_strings_lib, + dependencies: [absl_base_dep, absl_numeric_dep], +) + +absl_debugging_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_debugging_lib, + dependencies: [absl_base_dep, absl_strings_dep], +) + +absl_random_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_random_lib, + dependencies: [absl_base_dep, absl_strings_dep], +) + +absl_crc_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_crc_lib, + dependencies: [absl_base_dep], +) + +absl_time_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_time_lib, + dependencies: [absl_base_dep, absl_numeric_dep, absl_strings_dep], +) + +absl_types_dep = declare_dependency( + include_directories: absl_include_dir, +) + +absl_synchronization_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_synchronization_lib, + dependencies: [absl_base_dep, absl_debugging_dep, absl_time_dep], +) + +absl_container_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_container_lib, + dependencies: [ + absl_base_dep, + absl_debugging_dep, + absl_hash_dep, + absl_synchronization_dep, + absl_time_dep, + ], +) + +absl_flags_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_flags_lib, + dependencies: [ + absl_base_dep, + absl_container_dep, + absl_hash_dep, + absl_strings_dep, + absl_synchronization_dep, + ], +) + +absl_log_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_log_lib, + dependencies: [absl_base_dep, absl_strings_dep, absl_flags_dep], +) + +absl_status_dep = declare_dependency( + include_directories: absl_include_dir, + link_with: absl_status_lib, + dependencies: [absl_base_dep, absl_strings_dep], +) diff --git a/subprojects/packagefiles/gaviotatb/meson.build b/subprojects/packagefiles/gaviotatb/meson.build new file mode 100644 index 00000000..2a1b51bf --- /dev/null +++ b/subprojects/packagefiles/gaviotatb/meson.build @@ -0,0 +1,45 @@ +project('gaviotatb', 'c') + +gaviotatb_src = [ + 'gtb-probe.c', + 'gtb-dec.c', + 'gtb-att.c', + 'sysport/sysport.c', + 'compression/wrap.c', + 'compression/huffman/hzip.c', + 'compression/liblzf/lzf_c.c', + 'compression/liblzf/lzf_d.c', + 'compression/lzma/LzmaEnc.c', + 'compression/lzma/LzmaDec.c', + 'compression/lzma/Alloc.c', + 'compression/lzma/LzFind.c', + 'compression/lzma/Lzma86Enc.c', + 'compression/lzma/Lzma86Dec.c', + 'compression/lzma/Bra86.c' +] + +gaviotatb_includes = [ + '.', + 'sysport', + 'compression', + 'compression/huffman', + 'compression/liblzf', + 'compression/lzma', + 'compression/zlib' +] + +gaviota_lib = static_library('gaviota', + gaviotatb_src, + c_args : meson.get_compiler('c').get_supported_arguments([ + '-Dz_uLong=uLong', + '-Wno-misleading-indentation', + '-Wno-self-assign', + '-Wno-language-extension-token', + '-Wno-expansion-to-defined']), + include_directories : gaviotatb_includes) + +incdir = include_directories('.') + +gaviotatb_dep = declare_dependency( + link_with : gaviota_lib, + include_directories : gaviotatb_includes) diff --git a/tf/attention_policy_map.py b/tf/attention_policy_map.py new file mode 100644 index 00000000..7b3d988e --- /dev/null +++ b/tf/attention_policy_map.py @@ -0,0 +1,128 @@ +import numpy as np + + +move = np.arange(1, 8) + +diag = np.array([ + move + move*8, + move - move*8, + move*-1 - move*8, + move*-1 + move*8 +]) + +orthog = np.array([ + move, + move*-8, + move*-1, + move*8 +]) + +knight = np.array([ + [2 + 1*8], + [2 - 1*8], + [1 - 2*8], + [-1 - 2*8], + [-2 - 1*8], + [-2 + 1*8], + [-1 + 2*8], + [1 + 2*8] +]) + +promos = np.array([2*8, 3*8, 4*8]) +pawn_promotion = np.array([ + -1 + promos, + 0 + promos, + 1 + promos +]) + + +def make_map(): + """theoretically possible put-down squares (numpy array) for each pick-up square (list element). + squares are [0, 1, ..., 63] for [a1, b1, ..., h8]. squares after 63 are promotion squares. + each successive "row" beyond 63 (ie. 64:72, 72:80, 80:88) are for over-promotions to queen, rook, and bishop; + respectively. a pawn traverse to row 56:64 signifies a "default" promotion to a knight.""" + traversable = [] + for i in range(8): + for j in range(8): + sq = (8*i + j) + traversable.append( + sq + + np.sort( + np.int32( + np.concatenate(( + orthog[0][:7-j], orthog[2][:j], orthog[1][:i], orthog[3][:7-i], + diag[0][:np.min((7-i, 7-j))], diag[3][:np.min((7-i, j))], + diag[1][:np.min((i, 7-j))], diag[2][:np.min((i, j))], + knight[0] if i < 7 and j < 6 else [], knight[1] if i > 0 and j < 6 else [], + knight[2] if i > 1 and j < 7 else [], knight[3] if i > 1 and j > 0 else [], + knight[4] if i > 0 and j > 1 else [], knight[5] if i < 7 and j > 1 else [], + knight[6] if i < 6 and j > 0 else [], knight[7] if i < 6 and j < 7 else [], + pawn_promotion[0] if i == 6 and j > 0 else [], + pawn_promotion[1] if i == 6 else [], + pawn_promotion[2] if i == 6 and j < 7 else [], + )) + ) + ) + ) + z = np.zeros((64*64+8*24, 1858), dtype=np.int32) + # first loop for standard moves (for i in 0:1858, stride by 1) + i = 0 + for pickup_index, putdown_indices in enumerate(traversable): + for putdown_index in putdown_indices: + if putdown_index < 64: + z[putdown_index + (64*pickup_index), i] = 1 + i += 1 + # second loop for promotions (for i in 1792:1858, stride by ls[j]) + j = 0 + j1 = np.array([3, -2, 3, -2, 3]) + j2 = np.array([3, 3, -5, 3, 3, -5, 3, 3, 1]) + ls = np.append(j1, 1) + for k in range(6): + ls = np.append(ls, j2) + ls = np.append(ls, j1) + ls = np.append(ls, 0) + for pickup_index, putdown_indices in enumerate(traversable): + for putdown_index in putdown_indices: + if putdown_index >= 64: + pickup_file = pickup_index % 8 + promotion_file = putdown_index % 8 + promotion_rank = (putdown_index // 8) - 8 + z[4096 + pickup_file*24 + (promotion_file*3+promotion_rank), i] = 1 + i += ls[j] + j += 1 + + return z + +def make_pos_enc(): + traversable = [] + for i in range(8): + for j in range(8): + sq = (8*i + j) + traversable.append( + sq + + np.sort( + np.int32( + np.concatenate(( + orthog[0][:7-j], orthog[2][:j], orthog[1][:i], orthog[3][:7-i], + diag[0][:np.min((7-i, 7-j))], diag[3][:np.min((7-i, j))], + diag[1][:np.min((i, 7-j))], diag[2][:np.min((i, j))], + knight[0] if i < 7 and j < 6 else [], knight[1] if i > 0 and j < 6 else [], + knight[2] if i > 1 and j < 7 else [], knight[3] if i > 1 and j > 0 else [], + knight[4] if i > 0 and j > 1 else [], knight[5] if i < 7 and j > 1 else [], + knight[6] if i < 6 and j > 0 else [], knight[7] if i < 6 and j < 7 else [], + # pawn_promotion[0] if i == 6 and j > 0 else [], + # pawn_promotion[1] if i == 6 else [], + # pawn_promotion[2] if i == 6 and j < 7 else [], + )) + ) + ) + ) + + # pos_enc = np.zeros((1, 64, 88), dtype=np.float32) + pos_enc = np.zeros((1, 64, 64), dtype=np.float32) + for i, k in enumerate(traversable): + pos_enc[0][i][i] = -1. + for j in k: + pos_enc[0][i][j] = 1. + + return pos_enc diff --git a/tf/chunkparsefunc.py b/tf/chunkparsefunc.py new file mode 100644 index 00000000..3196ecb0 --- /dev/null +++ b/tf/chunkparsefunc.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 +# +# This file is part of Leela Chess. +# Copyright (C) 2021 Leela Chess Authors +# +# Leela Chess is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# Leela Chess is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with Leela Chess. If not, see . +import tensorflow as tf + + +def parse_function(planes, probs, winner, q, plies_left): + """ + Convert unpacked record batches to tensors for tensorflow training + """ + planes = tf.io.decode_raw(planes, tf.float32) + probs = tf.io.decode_raw(probs, tf.float32) + winner = tf.io.decode_raw(winner, tf.float32) + q = tf.io.decode_raw(q, tf.float32) + plies_left = tf.io.decode_raw(plies_left, tf.float32) + + planes = tf.reshape(planes, (-1, 112, 8, 8)) + probs = tf.reshape(probs, (-1, 1858)) + winner = tf.reshape(winner, (-1, 3)) + q = tf.reshape(q, (-1, 3)) + plies_left = tf.reshape(plies_left, (-1, 1)) + + return (planes, probs, winner, q, plies_left) diff --git a/tf/chunkparser.py b/tf/chunkparser.py index 981aed1e..9f6066e3 100644 --- a/tf/chunkparser.py +++ b/tf/chunkparser.py @@ -16,7 +16,6 @@ # # You should have received a copy of the GNU General Public License # along with Leela Chess. If not, see . - """ General comments on how chunkparser works. @@ -35,12 +34,13 @@ sample_record() also skips most training records to avoid sampling over-correlated positions since they typically are from sequential positions in a game. -Current implementation of "value focus" also is in sample_record() and works by +Current implementation of "diff focus" also is in sample_record() and works by probabilistically skipping records according to how accurate the no-search -eval ('orig_q') is compared to eval after search ('best_q'). It does not use -draw values at this point. Putting value focus here is efficient because -it runs in parallel workers and peeks at the records without requiring any -unpacking. +eval ('orig_q') is compared to eval after search ('best_q') as well as the +recorded policy_kld (a measure of difference between no search policy and the +final policy distribution). It does not use draw values at this point. Putting +diff focus here is efficient because it runs in parallel workers and peeks at +the records without requiring any unpacking. The constructor for chunkparser.ChunkParser() sets a bunch of class constants and creates a fixed number of parallel Python multiprocessing.Pipe objects, @@ -64,7 +64,6 @@ import random import shufflebuffer as sb import struct -import tensorflow as tf import unittest import gzip from select import select @@ -120,8 +119,6 @@ def chunk_reader(chunk_filenames, chunk_filename_queue): class ChunkParser: - # static batch size - BATCH_SIZE = 8 def __init__(self, chunks, @@ -130,16 +127,49 @@ def __init__(self, sample=1, buffer_size=1, batch_size=256, - value_focus_min=1, - value_focus_slope=0, + diff_focus_min=1, + diff_focus_slope=0, + diff_focus_q_weight=6.0, + diff_focus_pol_scale=3.5, workers=None): + self.inner = ChunkParserInner(self, chunks, expected_input_format, + shuffle_size, sample, buffer_size, + batch_size, diff_focus_min, + diff_focus_slope, diff_focus_q_weight, + diff_focus_pol_scale, workers) + + def shutdown(self): + """ + Terminates all the workers + """ + for i in range(len(self.processes)): + self.processes[i].terminate() + self.processes[i].join() + self.inner.readers[i].close() + self.inner.writers[i].close() + self.chunk_process.terminate() + self.chunk_process.join() + + def parse(self): + return self.inner.parse() + + def sequential(self): + return self.inner.sequential() + + +class ChunkParserInner: + def __init__(self, parent, chunks, expected_input_format, shuffle_size, + sample, buffer_size, batch_size, diff_focus_min, + diff_focus_slope, diff_focus_q_weight, diff_focus_pol_scale, + workers): """ Read data and yield batches of raw tensors. + 'parent' the outer chunk parser to store processes. Must not be stored by self directly or indirectly. 'chunks' list of chunk filenames. 'shuffle_size' is the size of the shuffle buffer. 'sample' is the rate to down-sample. - 'value_focus_min' and 'value_focus_slope' control value focus + 'diff_focus_min', 'diff_focus_slope', 'diff_focus_q_weight' and 'diff_focus_pol_scale' control diff focus 'workers' is the number of child workers to use. The data is represented in a number of formats through this dataflow @@ -166,9 +196,11 @@ def __init__(self, # set the down-sampling rate self.sample = sample - # set the min and slope for value focus, defaults accept all positions - self.value_focus_min = value_focus_min - self.value_focus_slope = value_focus_slope + # set the details for diff focus, defaults accept all positions + self.diff_focus_min = diff_focus_min + self.diff_focus_slope = diff_focus_slope + self.diff_focus_q_weight = diff_focus_q_weight + self.diff_focus_pol_scale = diff_focus_pol_scale # set the mini-batch size self.batch_size = batch_size # set number of elements in the shuffle buffer. @@ -177,43 +209,34 @@ def __init__(self, if workers is None: workers = max(1, mp.cpu_count() - 2) - print("Using {} worker processes.".format(workers)) - - # Start the child workers running - self.readers = [] - self.writers = [] - self.processes = [] - self.chunk_filename_queue = mp.Queue(maxsize=4096) - for _ in range(workers): - read, write = mp.Pipe(duplex=False) - p = mp.Process(target=self.task, - args=(self.chunk_filename_queue, write)) - p.daemon = True - self.processes.append(p) - p.start() - self.readers.append(read) - self.writers.append(write) - - self.chunk_process = mp.Process(target=chunk_reader, - args=(chunks, - self.chunk_filename_queue)) - self.chunk_process.daemon = True - self.chunk_process.start() + if workers > 0: + print("Using {} worker processes.".format(workers)) + + # Start the child workers running + self.readers = [] + self.writers = [] + parent.processes = [] + self.chunk_filename_queue = mp.Queue(maxsize=4096) + for _ in range(workers): + read, write = mp.Pipe(duplex=False) + p = mp.Process(target=self.task, + args=(self.chunk_filename_queue, write)) + p.daemon = True + parent.processes.append(p) + p.start() + self.readers.append(read) + self.writers.append(write) + + parent.chunk_process = mp.Process(target=chunk_reader, + args=(chunks, + self.chunk_filename_queue)) + parent.chunk_process.daemon = True + parent.chunk_process.start() + else: + self.chunks = chunks self.init_structs() - def shutdown(self): - """ - Terminates all the workers - """ - for i in range(len(self.readers)): - self.processes[i].terminate() - self.processes[i].join() - self.readers[i].close() - self.writers[i].close() - self.chunk_process.terminate() - self.chunk_process.join() - def init_structs(self): """ struct.Struct doesn't pickle, so it needs to be separately @@ -224,25 +247,6 @@ def init_structs(self): self.v4_struct = struct.Struct(V4_STRUCT_STRING) self.v3_struct = struct.Struct(V3_STRUCT_STRING) - @staticmethod - def parse_function(planes, probs, winner, q, plies_left): - """ - Convert unpacked record batches to tensors for tensorflow training - """ - planes = tf.io.decode_raw(planes, tf.float32) - probs = tf.io.decode_raw(probs, tf.float32) - winner = tf.io.decode_raw(winner, tf.float32) - q = tf.io.decode_raw(q, tf.float32) - plies_left = tf.io.decode_raw(plies_left, tf.float32) - - planes = tf.reshape(planes, (ChunkParser.BATCH_SIZE, 112, 8 * 8)) - probs = tf.reshape(probs, (ChunkParser.BATCH_SIZE, 1858)) - winner = tf.reshape(winner, (ChunkParser.BATCH_SIZE, 3)) - q = tf.reshape(q, (ChunkParser.BATCH_SIZE, 3)) - plies_left = tf.reshape(plies_left, (ChunkParser.BATCH_SIZE, )) - - return (planes, probs, winner, q, plies_left) - def convert_v6_to_tuple(self, content): """ Unpack a v6 binary record to 5-tuple (state, policy pi, result, q, m) @@ -294,11 +298,12 @@ def convert_v6_to_tuple(self, content): """ # unpack the V6 content from raw byte array, arbitrarily chose 4 2-byte values # for the 8 "reserved" bytes - (ver, input_format, probs, planes, us_ooo, us_oo, them_ooo, them_oo, stm, - rule50_count, invariance_info, dep_result, root_q, best_q, root_d, best_d, root_m, - best_m, plies_left, result_q, result_d, played_q, played_d, played_m, orig_q, - orig_d, orig_m, visits, played_idx, best_idx, reserved1, reserved2, reserved3, - reserved4) = self.v6_struct.unpack(content) + (ver, input_format, probs, planes, us_ooo, us_oo, them_ooo, them_oo, + stm, rule50_count, invariance_info, dep_result, root_q, best_q, + root_d, best_d, root_m, best_m, plies_left, result_q, result_d, + played_q, played_d, played_m, orig_q, orig_d, orig_m, visits, + played_idx, best_idx, reserved1, reserved2, reserved3, + reserved4) = self.v6_struct.unpack(content) """ v5 struct format was (8308 bytes total) int32 version (4 bytes) @@ -321,7 +326,7 @@ def convert_v6_to_tuple(self, content): float32 best_m (4 bytes) float32 plies_left (4 bytes) """ - # v3/4 data sometimes has a useful value in dep_ply_count (now invariance_info), + # v3/4 data sometimes has a useful value in dep_ply_count (now invariance_info), # so copy that over if the new ply_count is not populated. if plies_left == 0: plies_left = invariance_info @@ -370,7 +375,8 @@ def convert_v6_to_tuple(self, content): # Concatenate all byteplanes. Make the last plane all 1's so the NN can # detect edges of the board more easily aux_plus_6_plane = self.flat_planes[0] - if (input_format == 132 or input_format == 133) and invariance_info >= 128: + if (input_format == 132 + or input_format == 133) and invariance_info >= 128: aux_plus_6_plane = self.flat_planes[1] planes = planes.tobytes() + \ middle_planes + \ @@ -381,13 +387,13 @@ def convert_v6_to_tuple(self, content): assert len(planes) == ((8 * 13 * 1 + 8 * 1 * 1) * 8 * 8 * 4) if ver == V6_VERSION: - winner = struct.pack('fff', 0.5 * (1.0 - result_d + result_q), result_d, - 0.5 * (1.0 - result_d - result_q)) + winner = struct.pack('fff', 0.5 * (1.0 - result_d + result_q), + result_d, 0.5 * (1.0 - result_d - result_q)) else: dep_result = float(dep_result) assert dep_result == 1.0 or dep_result == -1.0 or dep_result == 0.0 winner = struct.pack('fff', dep_result == 1.0, dep_result == 0.0, - dep_result == -1.0) + dep_result == -1.0) best_q_w = 0.5 * (1.0 - best_d + best_q) best_q_l = 0.5 * (1.0 - best_d - best_q) @@ -396,12 +402,11 @@ def convert_v6_to_tuple(self, content): return (planes, probs, winner, best_q, plies_left) - def sample_record(self, chunkdata): """ Randomly sample through the v3/4/5/6 chunk data and select records in v6 format Downsampling to avoid highly correlated positions skips most records, and - value focus may also skip some records. + diff focus may also skip some records. """ version = chunkdata[0:4] if version == V6_VERSION: @@ -436,19 +441,63 @@ def sample_record(self, chunkdata): record += 48 * b'\x00' if version == V6_VERSION: - # value focus code, peek at best_q and orig_q from record (unpacks as tuple with one item) + # diff focus code, peek at best_q, orig_q and pol_kld from record (unpacks as tuple with one item) best_q = struct.unpack('f', record[8284:8288])[0] orig_q = struct.unpack('f', record[8328:8332])[0] - - # if orig_q is NaN, accept, else accept based on value focus - if not np.isnan(orig_q): + pol_kld = struct.unpack('f', record[8348:8352])[0] + + # if orig_q is NaN or pol_kld is 0, accept, else accept based on diff focus + if not np.isnan(orig_q) and pol_kld > 0: diff_q = abs(best_q - orig_q) - thresh_p = self.value_focus_min + self.value_focus_slope * diff_q + q_weight = self.diff_focus_q_weight + pol_scale = self.diff_focus_pol_scale + total = (q_weight * diff_q + pol_kld) / (q_weight + + pol_scale) + thresh_p = self.diff_focus_min + self.diff_focus_slope * total if thresh_p < 1.0 and random.random() > thresh_p: continue yield record + def single_file_gen(self, filename): + try: + with gzip.open(filename, 'rb') as chunk_file: + version = chunk_file.read(4) + chunk_file.seek(0) + if version == V6_VERSION: + record_size = self.v6_struct.size + elif version == V5_VERSION: + record_size = self.v5_struct.size + elif version == V4_VERSION: + record_size = self.v4_struct.size + elif version == V3_VERSION: + record_size = self.v3_struct.size + else: + print('Unknown version {} in file {}'.format( + version, filename)) + return + while True: + chunkdata = chunk_file.read(256 * record_size) + if len(chunkdata) == 0: + break + for item in self.sample_record(chunkdata): + yield item + + except: + print("failed to parse {}".format(filename)) + + def sequential_gen(self): + for filename in self.chunks: + for item in self.single_file_gen(filename): + yield item + + def sequential(self): + gen = self.sequential_gen() # read from all files in order in this process. + gen = self.tuple_gen(gen) # convert v6->tuple + gen = self.batch_gen(gen, allow_partial=False) # assemble into batches + for b in gen: + yield b + def task(self, chunk_filename_queue, writer): """ Run in fork'ed process, read data from chunkdatasrc, parsing, shuffling and @@ -457,32 +506,8 @@ def task(self, chunk_filename_queue, writer): self.init_structs() while True: filename = chunk_filename_queue.get() - try: - with gzip.open(filename, 'rb') as chunk_file: - version = chunk_file.read(4) - chunk_file.seek(0) - if version == V6_VERSION: - record_size = self.v6_struct.size - elif version == V5_VERSION: - record_size = self.v5_struct.size - elif version == V4_VERSION: - record_size = self.v4_struct.size - elif version == V3_VERSION: - record_size = self.v3_struct.size - else: - print('Unknown version {} in file {}'.format( - version, filename)) - continue - while True: - chunkdata = chunk_file.read(256 * record_size) - if len(chunkdata) == 0: - break - for item in self.sample_record(chunkdata): - writer.send_bytes(item) - - except: - print("failed to parse {}".format(filename)) - continue + for item in self.single_file_gen(filename): + writer.send_bytes(item) def v6_gen(self): """ @@ -516,7 +541,7 @@ def tuple_gen(self, gen): for r in gen: yield self.convert_v6_to_tuple(r) - def batch_gen(self, gen): + def batch_gen(self, gen, allow_partial=True): """ Pack multiple records into a single batch """ @@ -524,7 +549,7 @@ def batch_gen(self, gen): # a list because we need to reuse it. while True: s = list(itertools.islice(gen, self.batch_size)) - if not len(s): + if not len(s) or (not allow_partial and len(s) != self.batch_size): return yield (b''.join([x[0] for x in s]), b''.join([x[1] for x in s]), b''.join([x[2] for x in s]), b''.join([x[3] for x in s]), @@ -534,7 +559,7 @@ def parse(self): """ Read data from child workers and yield batches of unpacked records """ - gen = self.v6_gen() # read from workers + gen = self.v6_gen() # read from workers gen = self.tuple_gen(gen) # convert v6->tuple gen = self.batch_gen(gen) # assemble into batches for b in gen: @@ -635,54 +660,6 @@ def test_parsing(self): parser.shutdown() - def test_tensorflow_parsing(self): - """ - Test game position decoding pipeline including tensorflow. - """ - truth = self.generate_fake_pos() - batch_size = 4 - ChunkParser.BATCH_SIZE = batch_size - records = [] - for i in range(batch_size): - record = b'' - for j in range(2): - record += self.v4_record(*truth) - records.append(record) - - parser = ChunkParser(ChunkDataSrc(records), - shuffle_size=1, - workers=1, - batch_size=batch_size) - batchgen = parser.parse() - data = next(batchgen) - - planes = np.frombuffer(data[0], - dtype=np.float32, - count=112 * 8 * 8 * batch_size) - planes = planes.reshape(batch_size, 112, 8 * 8) - probs = np.frombuffer(data[1], - dtype=np.float32, - count=1858 * batch_size) - probs = probs.reshape(batch_size, 1858) - winner = np.frombuffer(data[2], dtype=np.float32, count=3 * batch_size) - winner = winner.reshape(batch_size, 3) - best_q = np.frombuffer(data[3], dtype=np.float32, count=3 * batch_size) - best_q = best_q.reshape(batch_size, 3) - - # Pass it through tensorflow - with tf.compat.v1.Session() as sess: - graph = ChunkParser.parse_function(data[0], data[1], data[2], - data[3]) - tf_planes, tf_probs, tf_winner, tf_q = sess.run(graph) - - for i in range(batch_size): - self.assertTrue((probs[i] == tf_probs[i]).all()) - self.assertTrue((planes[i] == tf_planes[i]).all()) - self.assertTrue((winner[i] == tf_winner[i]).all()) - self.assertTrue((best_q[i] == tf_q[i]).all()) - - parser.shutdown() - if __name__ == '__main__': unittest.main() diff --git a/tf/configs/example.yaml b/tf/configs/example.yaml index 40e56fae..35bf303f 100644 --- a/tf/configs/example.yaml +++ b/tf/configs/example.yaml @@ -29,9 +29,20 @@ training: - 130000 policy_loss_weight: 1.0 # weight of policy loss value_loss_weight: 1.0 # weight of value loss + moves_left_loss_weight: 1.0 # weight of moves-left loss path: '/path/to/store/networks' # network storage dir model: filters: 64 residual_blocks: 6 + se_ratio: 2 # Squeeze Excite structural network architecture. + policy: 'attention' # attention policy fields: + pol_embedding_size: 64 # embedding vector size + pol_encoder_layers: 1 # number of intermediate attention layers in the policy head + pol_encoder_heads: 4 # number of attention heads in encoder layers + pol_encoder_d_model: 64 # size of the Q, K, & V vectors in encoder layers -- divisible by encoder_heads + pol_encoder_dff: 128 # size of the largest dense layer in encoder block feed-forward network + policy_d_model: 64 # size of the query and key vectors in final attention layer + value: 'wdl' + moves_left: 'v1' ... diff --git a/tf/make_model.py b/tf/make_model.py new file mode 100644 index 00000000..22507bdf --- /dev/null +++ b/tf/make_model.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 +import argparse +import os +import yaml +import tfprocess + +argparser = argparse.ArgumentParser(description='Convert net to model.') +argparser.add_argument('--start', + type=int, + default=0, + help='Offset to set global_step to.') +argparser.add_argument('--cfg', + type=argparse.FileType('r'), + help='yaml configuration with training parameters') +args = argparser.parse_args() +cfg = yaml.safe_load(args.cfg.read()) +print(yaml.dump(cfg, default_flow_style=False)) +START_FROM = args.start + +tfp = tfprocess.TFProcess(cfg) +tfp.init_net() +tfp.global_step.assign(START_FROM) + +root_dir = os.path.join(cfg['training']['path'], cfg['name']) +if not os.path.exists(root_dir): + os.makedirs(root_dir) +tfp.manager.save(checkpoint_number=START_FROM) +print("Wrote model to {}".format(tfp.manager.latest_checkpoint)) +path = os.path.join(tfp.root_dir, tfp.cfg['name']) +leela_path = path + "-" + str(START_FROM) +swa_path = path + "-swa-" + str(START_FROM) +tfp.net.pb.training_params.training_steps = START_FROM +tfp.save_leelaz_weights(leela_path) +if tfp.swa_enabled: + tfp.save_swa_weights(swa_path) + diff --git a/tf/model_to_net.py b/tf/model_to_net.py new file mode 100644 index 00000000..c894f366 --- /dev/null +++ b/tf/model_to_net.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +import argparse +import os +import yaml +import tfprocess + +argparser = argparse.ArgumentParser(description='Convert model to net.') +argparser.add_argument('--cfg', + type=argparse.FileType('r'), + help='yaml configuration with training parameters') +args = argparser.parse_args() +cfg = yaml.safe_load(args.cfg.read()) +print(yaml.dump(cfg, default_flow_style=False)) + +tfp = tfprocess.TFProcess(cfg) +tfp.init_net() + +tfp.restore() + +root_dir = os.path.join(cfg['training']['path'], cfg['name']) +if not os.path.exists(root_dir): + os.makedirs(root_dir) +path = os.path.join(tfp.root_dir, tfp.cfg['name']) +steps = tfp.global_step.read_value().numpy() +leela_path = path + "-" + str(steps) +swa_path = path + "-swa-" + str(steps) +tfp.net.pb.training_params.training_steps = steps +tfp.save_leelaz_weights(leela_path) +if tfp.swa_enabled: + tfp.save_swa_weights(swa_path) + diff --git a/tf/net.py b/tf/net.py index de779ca0..7a51dd8f 100755 --- a/tf/net.py +++ b/tf/net.py @@ -11,6 +11,8 @@ LC0_MINOR_WITH_INPUT_TYPE_3 = 25 LC0_MINOR_WITH_INPUT_TYPE_4 = 26 LC0_MINOR_WITH_INPUT_TYPE_5 = 27 +LC0_MINOR_WITH_MISH = 29 +LC0_MINOR_WITH_ATTN_BODY = 30 LC0_PATCH = 0 WEIGHTS_MAGIC = 0x1c0 @@ -23,6 +25,7 @@ def nested_getattr(obj, attr): class Net: + def __init__(self, net=pb.NetworkFormat.NETWORK_SE_WITH_HEADFORMAT, input=pb.NetworkFormat.INPUT_CLASSICAL_112_PLANE, @@ -49,13 +52,23 @@ def __init__(self, self.set_policyformat(policy) self.set_valueformat(value) self.set_movesleftformat(moves_left) + self.set_defaultactivation(pb.NetworkFormat.DEFAULT_ACTIVATION_RELU) def set_networkformat(self, net): self.pb.format.network_format.network = net + if net == pb.NetworkFormat.NETWORK_ATTENTIONBODY_WITH_HEADFORMAT \ + and self.pb.min_version.minor < LC0_MINOR_WITH_ATTN_BODY: + self.pb.min_version.minor = LC0_MINOR_WITH_ATTN_BODY def set_policyformat(self, policy): self.pb.format.network_format.policy = policy + def set_headcount(self, headcount): + self.pb.weights.headcount = headcount + + def set_pol_headcount(self, headcount): + self.pb.weights.pol_headcount = headcount + def set_valueformat(self, value): self.pb.format.network_format.value = value @@ -78,6 +91,46 @@ def set_input(self, input_format): elif input_format != pb.NetworkFormat.INPUT_CLASSICAL_112_PLANE: self.pb.min_version.minor = LC0_MINOR_WITH_INPUT_TYPE_3 + def set_defaultactivation(self, activation): + self.pb.format.network_format.default_activation = activation + if activation == pb.NetworkFormat.DEFAULT_ACTIVATION_MISH: + if self.pb.min_version.minor < LC0_MINOR_WITH_MISH: + self.pb.min_version.minor = LC0_MINOR_WITH_MISH + + def set_smolgen_activation(self, activation): + self.pb.format.network_format.smolgen_activation = activation + if self.pb.min_version.minor < LC0_MINOR_WITH_ATTN_BODY: + self.pb.min_version.minor = LC0_MINOR_WITH_ATTN_BODY + return None + + def set_ffn_activation(self, activation): + self.pb.format.network_format.ffn_activation = activation + if self.pb.min_version.minor < LC0_MINOR_WITH_ATTN_BODY: + self.pb.min_version.minor = LC0_MINOR_WITH_ATTN_BODY + return None + + def activation(self, name): + if name == "relu": + return pb.NetworkFormat.ACTIVATION_RELU + elif name == "tanh": + return pb.NetworkFormat.ACTIVATION_TANH + elif name == "sigmoid": + return pb.NetworkFormat.ACTIVATION_SIGMOID + elif name == "softmax": + return pb.NetworkFormat.ACTIVATION_SOFTMAX + elif name == "selu": + return pb.NetworkFormat.ACTIVATION_SELU + elif name == "mish": + return pb.NetworkFormat.ACTIVATION_MISH + elif name == "swish": + return pb.NetworkFormat.ACTIVATION_SWISH + elif name == "relu_2" or name == "sqrrelu": + return pb.NetworkFormat.ACTIVATION_RELU_2 + elif name == "default": + return pb.NetworkFormat.ACTIVATION_DEFAULT + else: + return pb.NetworkFormat.ACTIVATION_NONE + def get_weight_amounts(self): value_weights = 8 policy_weights = 6 @@ -226,6 +279,7 @@ def save_proto(self, filename): def tf_name_to_pb_name(self, name): """Given Tensorflow variable name returns the protobuf name and index of residual block if weight belong in a residual block.""" + def convblock_to_bp(w): w = w.split(':')[0] d = { @@ -253,7 +307,9 @@ def se_to_bp(l, w): return d[w] + str(n) def value_to_bp(l, w): - if l == 'dense1': + if l == 'embedding': + n = '' + elif l == 'dense1': n = 1 elif l == 'dense2': n = 2 @@ -265,14 +321,85 @@ def value_to_bp(l, w): return d[w].format(n) - def policy_to_bp(w): + def conv_policy_to_bp(w): w = w.split(':')[0] d = {'kernel': 'ip_pol_w', 'bias': 'ip_pol_b'} return d[w] + def attn_pol_to_bp(l, w): + if l == 'wq': + n = 2 + elif l == 'wk': + n = 3 + elif l == 'ppo': + n = 4 + else: + raise ValueError( + 'Unable to decode attn_policy weight {}/{}'.format(l, w)) + w = w.split(':')[0] + d = {'kernel': 'ip{}_pol_w', 'bias': 'ip{}_pol_b'} + + return d[w].format(n) + + def encoder_to_bp(l, w): + if l == 'ln1': + n = 1 + elif l == 'ln2': + n = 2 + else: + raise ValueError( + 'Unable to decode encoder weight {}/{}'.format(l, w)) + w = w.split(':')[0] + d = {'gamma': 'ln{}_gammas', 'beta': 'ln{}_betas'} + + return d[w].format(n) + + def mha_to_bp(l, w): + s = '' + if l.startswith('dense'): + s = 'dense' + elif l.startswith('w'): + s = l[1] + else: + raise ValueError('Unable to decode mha weight {}/{}'.format( + l, w)) + w = w.split(':')[0] + d = {'kernel': '{}_w', 'bias': '{}_b'} + + return d[w].format(s) + + def mha_smolgen_to_bp(l, w): + s = { + 'compress': 'compress', + 'hidden1_dense': 'dense1_{}', + 'hidden1_ln': 'ln1_{}', + 'gen_from': 'dense2_{}', + 'gen_from_ln': 'ln2_{}' + } + if s[l] is None: + raise ValueError( + 'Unable to decode mha smolgen weight {}/{}'.format(l, w)) + w = w.split(':')[0] + d = { + 'kernel': 'w', + 'bias': 'b', + 'gamma': 'gammas', + 'beta': 'betas' + } + + return s[l].format(d[w]) + + def ffn_to_bp(l, w): + w = w.split(':')[0] + d = {'kernel': '{}_w', 'bias': '{}_b'} + + return d[w].format(l) + def moves_left_to_bp(l, w): - if l == 'dense1': + if l == 'embedding': + n = '' + elif l == 'dense1': n = 1 elif l == 'dense2': n = 2 @@ -289,6 +416,8 @@ def moves_left_to_bp(l, w): weights_name = layers[-1] pb_name = None block = None + encoder_block = None + pol_encoder_block = None if base_layer == 'input': pb_name = 'input.' + convblock_to_bp(weights_name) @@ -296,16 +425,31 @@ def moves_left_to_bp(l, w): pb_name = 'policy1.' + convblock_to_bp(weights_name) elif base_layer == 'policy': if 'dense' in layers[1]: - pb_name = policy_to_bp(weights_name) + pb_name = conv_policy_to_bp(weights_name) + elif layers[1] == 'embedding': + if layers[2].split(':')[0] == 'kernel': + pb_name = 'ip_pol_w' + else: + pb_name = 'ip_pol_b' + elif layers[1] == 'attention': + pb_name = attn_pol_to_bp(layers[2], weights_name) + elif layers[1].startswith('enc_layer_'): + pol_encoder_block = int(layers[1].split('_')[2]) - 1 + if layers[2] == 'mha': + pb_name = 'mha.' + mha_to_bp(layers[3], weights_name) + elif layers[2] == 'ffn': + pb_name = 'ffn.' + ffn_to_bp(layers[3], weights_name) + else: + pb_name = encoder_to_bp(layers[2], weights_name) else: pb_name = 'policy.' + convblock_to_bp(weights_name) elif base_layer == 'value': - if 'dense' in layers[1]: + if 'dense' in layers[1] or 'embedding' in layers[1]: pb_name = value_to_bp(layers[1], weights_name) else: pb_name = 'value.' + convblock_to_bp(weights_name) elif base_layer == 'moves_left': - if 'dense' in layers[1]: + if 'dense' in layers[1] or 'embedding' in layers[1]: pb_name = moves_left_to_bp(layers[1], weights_name) else: pb_name = 'moves_left.' + convblock_to_bp(weights_name) @@ -317,8 +461,33 @@ def moves_left_to_bp(l, w): pb_name = 'conv2.' + convblock_to_bp(weights_name) elif layers[1] == 'se': pb_name = 'se.' + se_to_bp(layers[-2], weights_name) + elif base_layer.startswith('encoder'): + encoder_block = int(base_layer.split('_')[1]) - 1 + if layers[1] == 'mha': + if layers[2] == 'smolgen': + pb_name = 'mha.smolgen.' + mha_smolgen_to_bp( + layers[3], weights_name) + else: + pb_name = 'mha.' + mha_to_bp(layers[2], weights_name) + elif layers[1] == 'ffn': + pb_name = 'ffn.' + ffn_to_bp(layers[2], weights_name) + else: + pb_name = encoder_to_bp(layers[1], weights_name) + elif base_layer == 'embedding': + if layers[1] == 'mult_gate' or layers[1] == 'add_gate': + if layers[2].split(':')[0] == 'gate': + pb_name = 'ip_{}'.format(layers[1]) + elif layers[1].split(':')[0] == 'kernel': + pb_name = 'ip_emb_w' + elif layers[1].split(':')[0] == 'bias': + pb_name = 'ip_emb_b' + elif base_layer == 'smol_weight_gen': + if layers[1].split(':')[0] == 'kernel': + pb_name = 'smolgen_w' + else: + pb_name = 'smolgen_b' - return (pb_name, block) + return (pb_name, block, pol_encoder_block, encoder_block) def get_weights_v2(self, names): # `names` is a list of Tensorflow tensor names to get from the protobuf. @@ -333,16 +502,25 @@ def get_weights_v2(self, names): if 'renorm' in name: # Renorm variables are not populated. continue + if 'headcount' in tf_name: + # headcount is set with set_headcount() + continue - pb_name, block = self.tf_name_to_pb_name(name) + pb_name, block, pol_encoder_block, encoder_block = self.tf_name_to_pb_name( + name) if pb_name is None: raise ValueError( "Don't know where to store weight in protobuf: {}".format( name)) - if block == None: - pb_weights = self.pb.weights + if block is None: + if pol_encoder_block is not None: + pb_weights = self.pb.weights.pol_encoder[pol_encoder_block] + elif encoder_block is not None: + pb_weights = self.pb.weights.encoder[encoder_block] + else: + pb_weights = self.pb.weights else: pb_weights = self.pb.weights.residual[block] @@ -474,19 +652,35 @@ def fill_net_v2(self, all_weights): weights = np.square(weights) - 1e-5 name = name.replace('stddev', 'variance') - if name == 'input/conv2d/kernel:0' and self.pb.format.network_format.input < pb.NetworkFormat.INPUT_112_WITH_CANONICALIZATION_HECTOPLIES: - # 50 move rule is the 110th input, or 109 starting from 0. - weights[:, 109, :, :] /= 99 + if self.pb.format.network_format.input < pb.NetworkFormat.INPUT_112_WITH_CANONICALIZATION_HECTOPLIES: + if name == 'input/conv2d/kernel:0': + # 50 move rule is the 110th input, or 109 starting from 0. + weights[:, 109, :, :] /= 99 + elif name == 'embedding/kernel:0': + weights[:, 109] /= 99 - pb_name, block = self.tf_name_to_pb_name(name) + pb_name, block, pol_encoder_block, encoder_block = self.tf_name_to_pb_name( + name) if pb_name is None: raise ValueError( "Don't know where to store weight in protobuf: {}".format( name)) - if block == None: - pb_weights = self.pb.weights + if block is None: + if pol_encoder_block is not None: + assert pol_encoder_block >= 0 + while pol_encoder_block >= len( + self.pb.weights.pol_encoder): + self.pb.weights.pol_encoder.add() + pb_weights = self.pb.weights.pol_encoder[pol_encoder_block] + elif encoder_block is not None: + assert encoder_block >= 0 + while encoder_block >= len(self.pb.weights.encoder): + self.pb.weights.encoder.add() + pb_weights = self.pb.weights.encoder[encoder_block] + else: + pb_weights = self.pb.weights else: assert block >= 0 while block >= len(self.pb.weights.residual): diff --git a/tf/net_to_model.py b/tf/net_to_model.py index c8a19d82..02772f82 100755 --- a/tf/net_to_model.py +++ b/tf/net_to_model.py @@ -3,7 +3,6 @@ import os import yaml import tfprocess -from net import Net argparser = argparse.ArgumentParser(description='Convert net to model.') argparser.add_argument('net', @@ -26,8 +25,8 @@ START_FROM = args.start tfp = tfprocess.TFProcess(cfg) -tfp.init_net_v2() -tfp.replace_weights_v2(args.net, args.ignore_errors) +tfp.init_net() +tfp.replace_weights(args.net, args.ignore_errors) tfp.global_step.assign(START_FROM) root_dir = os.path.join(cfg['training']['path'], cfg['name']) diff --git a/tf/requirements.txt b/tf/requirements.txt index dbadfd70..4718391c 100644 --- a/tf/requirements.txt +++ b/tf/requirements.txt @@ -1,4 +1,4 @@ numpy==1.13.3 -tensorflow==2.0.3 +tensorflow==2.5.1 tensorflow-tensorboard==0.4.0rc2 protobuf==3.12.1 diff --git a/tf/start.sh b/tf/start.sh index 42496776..37d1179e 100755 --- a/tf/start.sh +++ b/tf/start.sh @@ -7,7 +7,9 @@ REPO="origin" ROOT="/work/lc0" NETDIR="$ROOT/networks/upload" GAMEFILE="$HOME/.lc0.dat" +LATESTFILE="$HOME/.lc0.latest.dat" RAMDISK="/ramdisk" +MIN_GAP=10 function usage() { @@ -70,11 +72,21 @@ train() { mv -v $2.pb.gz $NETDIR } +delay_count=$((MIN_GAP+1)) while true do - if [ -f "$ROOT/data/$file" ] + latest_num=0 + if [ -f "$LATESTFILE" ] then + latest_num=$(cat $LATESTFILE) + fi + if [[ $delay_count -gt $MIN_GAP && ( $latest_num -gt $game_num || -f "$ROOT/data-rescored/$file" ) ]] + then + if [ $latest_num -gt $game_num ] + then + game_num=$latest_num + fi echo "" # prepare ramdisk @@ -95,8 +107,10 @@ do game_num=$((game_num + GAMES)) file="training.${game_num}.gz" echo "Waiting for '$file'" + delay_count=1 else echo -n "." sleep 60 + delay_count=$((delay_count+1)) fi done diff --git a/tf/tfprocess.py b/tf/tfprocess.py index a290ce24..17d198fe 100644 --- a/tf/tfprocess.py +++ b/tf/tfprocess.py @@ -23,6 +23,7 @@ import time import bisect import lc0_az_policy_map +import attention_policy_map as apm import proto.net_pb2 as pb from functools import reduce import operator @@ -30,7 +31,41 @@ from net import Net +def square_relu(x): + return tf.nn.relu(x)**2 + + +class Gating(tf.keras.layers.Layer): + + def __init__(self, name=None, additive=True, init_value=None, **kwargs): + self.additive = additive + if init_value is None: + init_value = 0 if self.additive else 1 + self.init_value = init_value + super().__init__(name=name, **kwargs) + + def build(self, input_shape): + self.gate = self.add_weight(name='gate', + shape=input_shape[1:], + constraint=tf.keras.constraints.NonNeg() + if not self.additive else None, + initializer=tf.constant_initializer( + self.init_value), + trainable=True) + + def call(self, inputs): + return tf.add(inputs, self.gate) if self.additive else tf.multiply( + inputs, self.gate) + + +def ma_gating(inputs, name): + out = Gating(name=name + '/mult_gate', additive=False)(inputs) + out = Gating(name=name + '/add_gate', additive=True)(out) + return out + + class ApplySqueezeExcitation(tf.keras.layers.Layer): + def __init__(self, **kwargs): super(ApplySqueezeExcitation, self).__init__(**kwargs) @@ -48,6 +83,7 @@ def call(self, inputs): class ApplyPolicyMap(tf.keras.layers.Layer): + def __init__(self, **kwargs): super(ApplyPolicyMap, self).__init__(**kwargs) self.fc1 = tf.constant(lc0_az_policy_map.make_map()) @@ -58,7 +94,58 @@ def call(self, inputs): tf.cast(self.fc1, h_conv_pol_flat.dtype)) +class ApplyAttentionPolicyMap(tf.keras.layers.Layer): + + def __init__(self, **kwargs): + super(ApplyAttentionPolicyMap, self).__init__(**kwargs) + self.fc1 = tf.constant(apm.make_map()) + + def call(self, logits, pp_logits): + logits = tf.concat([ + tf.reshape(logits, [-1, 64 * 64]), + tf.reshape(pp_logits, [-1, 8 * 24]) + ], + axis=1) + return tf.matmul(logits, tf.cast(self.fc1, logits.dtype)) + + +class Metric: + + def __init__(self, short_name, long_name, suffix='', **kwargs): + self.short_name = short_name + self.long_name = long_name + self.suffix = suffix + self.value = 0.0 + self.count = 0 + + def assign(self, value): + self.value = value + self.count = 1 + + def accumulate(self, value): + if self.count > 0: + self.value = self.value + value + self.count = self.count + 1 + else: + self.assign(value) + + def merge(self, other): + assert self.short_name == other.short_name + self.value = self.value + other.value + self.count = self.count + other.count + + def get(self): + if self.count == 0: + return self.value + return self.value / self.count + + def reset(self): + self.value = 0.0 + self.count = 0 + + class TFProcess: + def __init__(self, cfg): self.cfg = cfg self.net = Net() @@ -66,10 +153,60 @@ def __init__(self, cfg): self.cfg['name']) # Network structure - self.RESIDUAL_FILTERS = self.cfg['model']['filters'] - self.RESIDUAL_BLOCKS = self.cfg['model']['residual_blocks'] - self.SE_ratio = self.cfg['model']['se_ratio'] + self.RESIDUAL_FILTERS = self.cfg['model'].get('filters', 0) + self.RESIDUAL_BLOCKS = self.cfg['model'].get('residual_blocks', 0) + self.SE_ratio = self.cfg['model'].get('se_ratio', 0) + self.encoder_layers = self.cfg['model'].get('encoder_layers', 0) + self.encoder_heads = self.cfg['model'].get('encoder_heads', 2) + assert (self.RESIDUAL_BLOCKS > 0) != (self.encoder_layers > 0), \ + "Nets with both encoder layers and residual blocks are not supported" + if self.encoder_layers > 0: + self.RESIDUAL_FILTERS = self.cfg['model']['embedding_size'] + self.embedding_size = self.RESIDUAL_FILTERS + self.policy_channels = self.cfg['model'].get('policy_channels', 32) + self.pol_embedding_size = self.cfg['model'].get( + 'pol_embedding_size', self.RESIDUAL_FILTERS) + self.val_embedding_size = self.cfg['model'].get( + 'value_embedding_size', 32) + self.mov_embedding_size = self.cfg['model'].get( + 'moves_left_embedding_size', 8) + #policy head + self.pol_encoder_layers = (0 if self.encoder_layers > 0 else 1) + #logic is to explictly warn users who set both in yaml + if self.cfg['model'].get('pol_encoder_layers') is not None: + self.pol_encoder_layers = self.cfg['model'].get( + 'pol_encoder_layers') + assert not ((self.pol_encoder_layers > 0) and (self.encoder_layers > 0)), \ + "Nets with both body encoder layers and policy encoder layers are not supported" + self.pol_encoder_heads = self.cfg['model'].get('pol_encoder_heads', 2) + self.pol_encoder_d_model = self.cfg['model'].get( + 'pol_encoder_d_model', self.RESIDUAL_FILTERS) + self.pol_encoder_dff = self.cfg['model'].get( + 'pol_encoder_dff', (self.RESIDUAL_FILTERS * 1.5) // 1) + self.policy_d_model = self.cfg['model'].get('policy_d_model', + self.RESIDUAL_FILTERS) + + #encoder body + self.input_gate = self.cfg['model'].get('input_gate') + self.encoder_d_model = self.cfg['model'].get('encoder_d_model') + self.encoder_dff = self.cfg['model'].get( + 'encoder_dff', (self.RESIDUAL_FILTERS * 1.5) // 1) + self.policy_d_model = self.cfg['model'].get('policy_d_model', + self.RESIDUAL_FILTERS) + self.arc_encoding = self.cfg['model'].get('arc_encoding', True) + self.square_relu_ffn = self.cfg['model'].get('square_relu_ffn', False) + + self.use_smolgen = self.cfg['model'].get('use_smolgen', False) + self.smolgen_hidden_channels = self.cfg['model'].get( + 'smolgen_hidden_channels', 16) + self.smolgen_hidden_sz = self.cfg['model'].get('smolgen_hidden_sz', + 128) + self.smolgen_gen_sz = self.cfg['model'].get('smolgen_gen_sz', 128) + self.smolgen_activation = self.cfg['model'].get( + 'smolgen_activation', 'swish') + + self.dropout_rate = self.cfg['model'].get('dropout_rate', 0.0) precision = self.cfg['training'].get('precision', 'single') loss_scale = self.cfg['training'].get('loss_scale', 128) self.virtual_batch_size = self.cfg['model'].get( @@ -89,16 +226,23 @@ def __init__(self, cfg): value_head = self.cfg['model'].get('value', 'wdl') moves_left_head = self.cfg['model'].get('moves_left', 'v1') input_mode = self.cfg['model'].get('input_type', 'classic') + default_activation = self.cfg['model'].get('default_activation', + 'relu') self.POLICY_HEAD = None self.VALUE_HEAD = None self.MOVES_LEFT_HEAD = None self.INPUT_MODE = None + self.DEFAULT_ACTIVATION = None if policy_head == "classical": self.POLICY_HEAD = pb.NetworkFormat.POLICY_CLASSICAL elif policy_head == "convolution": self.POLICY_HEAD = pb.NetworkFormat.POLICY_CONVOLUTION + elif policy_head == "attention": + self.POLICY_HEAD = pb.NetworkFormat.POLICY_ATTENTION + if self.pol_encoder_layers > 0: + self.net.set_pol_headcount(self.pol_encoder_heads) else: raise ValueError( "Unknown policy head format: {}".format(policy_head)) @@ -149,6 +293,32 @@ def __init__(self, cfg): self.net.set_input(self.INPUT_MODE) + if default_activation == "relu": + self.net.set_defaultactivation( + pb.NetworkFormat.DEFAULT_ACTIVATION_RELU) + self.DEFAULT_ACTIVATION = 'relu' + elif default_activation == "mish": + self.net.set_defaultactivation( + pb.NetworkFormat.DEFAULT_ACTIVATION_MISH) + try: + self.DEFAULT_ACTIVATION = tf.keras.activations.mish + except AttributeError: + import tensorflow_addons as tfa + self.DEFAULT_ACTIVATION = tfa.activations.mish + else: + raise ValueError("Unknown default activation type: {}".format( + default_activation)) + + if self.encoder_layers > 0: + self.net.set_headcount(self.encoder_heads) + self.net.set_networkformat( + pb.NetworkFormat.NETWORK_ATTENTIONBODY_WITH_HEADFORMAT) + self.net.set_smolgen_activation( + self.net.activation(self.smolgen_activation)) + self.net.set_ffn_activation( + self.net.activation( + 'sqrrelu' if self.square_relu_ffn else 'default')) + self.swa_enabled = self.cfg['training'].get('swa', False) # Limit momentum of SWA exponential average to 1 - 1/(swa_max_n + 1) @@ -160,10 +330,20 @@ def __init__(self, cfg): self.renorm_momentum = self.cfg['training'].get( 'renorm_momentum', 0.99) - gpus = tf.config.experimental.list_physical_devices('GPU') - tf.config.experimental.set_visible_devices(gpus[self.cfg['gpu']], - 'GPU') - tf.config.experimental.set_memory_growth(gpus[self.cfg['gpu']], True) + if self.cfg['gpu'] == 'all': + gpus = tf.config.experimental.list_physical_devices('GPU') + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + self.strategy = tf.distribute.MirroredStrategy() + tf.distribute.experimental_set_strategy(self.strategy) + else: + gpus = tf.config.experimental.list_physical_devices('GPU') + print(gpus) + tf.config.experimental.set_visible_devices(gpus[self.cfg['gpu']], + 'GPU') + tf.config.experimental.set_memory_growth(gpus[self.cfg['gpu']], + True) + self.strategy = None if self.model_dtype == tf.float16: tf.keras.mixed_precision.experimental.set_policy('mixed_float16') @@ -172,26 +352,38 @@ def __init__(self, cfg): trainable=False, dtype=tf.int64) - def init_v2(self, train_dataset, test_dataset, validation_dataset=None): - self.train_dataset = train_dataset - self.train_iter = iter(train_dataset) - self.test_dataset = test_dataset - self.test_iter = iter(test_dataset) - self.validation_dataset = validation_dataset - self.init_net_v2() + def init(self, train_dataset, test_dataset, validation_dataset=None): + if self.strategy is not None: + self.train_dataset = self.strategy.experimental_distribute_dataset( + train_dataset) + else: + self.train_dataset = train_dataset + self.train_iter = iter(self.train_dataset) + if self.strategy is not None: + self.test_dataset = self.strategy.experimental_distribute_dataset( + test_dataset) + else: + self.test_dataset = test_dataset + self.test_iter = iter(self.test_dataset) + if self.strategy is not None and validation_dataset is not None: + self.validation_dataset = self.strategy.experimental_distribute_dataset( + validation_dataset) + else: + self.validation_dataset = validation_dataset + if self.strategy is not None: + this = self + with self.strategy.scope(): + this.init_net() + else: + self.init_net() - def init_net_v2(self): + def init_net(self): self.l2reg = tf.keras.regularizers.l2(l=0.5 * (0.0001)) - input_var = tf.keras.Input(shape=(112, 8 * 8)) - x_planes = tf.keras.layers.Reshape([112, 8, 8])(input_var) - policy, value, moves_left = self.construct_net_v2(x_planes) - if self.moves_left: - outputs = [policy, value, moves_left] - else: - outputs = [policy, value] + input_var = tf.keras.Input(shape=(112, 8, 8)) + outputs = self.construct_net(input_var) self.model = tf.keras.Model(inputs=input_var, outputs=outputs) - # swa_count initialized reguardless to make checkpoint code simpler. + # swa_count initialized regardless to make checkpoint code simpler. self.swa_count = tf.Variable(0., name='swa_count', trainable=False) self.swa_weights = None if self.swa_enabled: @@ -200,13 +392,36 @@ def init_net_v2(self): tf.Variable(w, trainable=False) for w in self.model.weights ] - self.active_lr = 0.01 - self.optimizer = tf.keras.optimizers.SGD( - learning_rate=lambda: self.active_lr, momentum=0.9, nesterov=True) + self.active_lr = tf.Variable(0.01, trainable=False) + # All 'new' (TF 2.10 or newer non-legacy) optimizers must have learning_rate updated manually. + self.update_lr_manually = False + # Be sure not to set new_optimizer before TF 2.11, or unless you edit the code to specify a new optimizer explicitly. + if self.cfg['training'].get('new_optimizer'): + self.optimizer = tf.keras.optimizers.SGD( + learning_rate=self.active_lr, momentum=0.9, nesterov=True) + self.update_lr_manually = True + else: + try: + self.optimizer = tf.keras.optimizers.legacy.SGD( + learning_rate=lambda: self.active_lr, + momentum=0.9, + nesterov=True) + except AttributeError: + self.optimizer = tf.keras.optimizers.SGD( + learning_rate=lambda: self.active_lr, + momentum=0.9, + nesterov=True) self.orig_optimizer = self.optimizer + try: + self.aggregator = self.orig_optimizer.aggregate_gradients + except AttributeError: + self.aggregator = self.orig_optimizer.gradient_aggregator if self.loss_scale != 1: self.optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer( self.optimizer, self.loss_scale) + if self.cfg['training'].get('lookahead_optimizer'): + import tensorflow_addons as tfa + self.optimizer = tfa.optimizers.Lookahead(self.optimizer) def correct_policy(target, output): output = tf.cast(output, tf.float32) @@ -238,8 +453,6 @@ def policy_accuracy(target, output): self.policy_accuracy_fn = policy_accuracy - self.policy_accuracy_fn = policy_accuracy - def moves_left_mean_error_fn(target, output): output = tf.cast(output, tf.float32) return tf.reduce_mean(tf.abs(target - output)) @@ -317,7 +530,11 @@ def moves_left_loss(target, output): scale = 20.0 target = target / scale output = tf.cast(output, tf.float32) / scale - huber = tf.keras.losses.Huber(10.0 / scale) + if self.strategy is not None: + huber = tf.keras.losses.Huber( + 10.0 / scale, reduction=tf.keras.losses.Reduction.NONE) + else: + huber = tf.keras.losses.Huber(10.0 / scale) return tf.reduce_mean(huber(target, output)) else: moves_left_loss = None @@ -331,9 +548,10 @@ def moves_left_loss(target, output): moves_loss_w = self.cfg['training']['moves_left_loss_weight'] else: moves_loss_w = tf.constant(0.0, dtype=tf.float32) + reg_term_w = self.cfg['training'].get('reg_term_weight', 1.0) - def _lossMix(policy, value, moves_left): - return pol_loss_w * policy + val_loss_w * value + moves_loss_w * moves_left + def _lossMix(policy, value, moves_left, reg_term): + return pol_loss_w * policy + val_loss_w * value + moves_loss_w * moves_left + reg_term_w * reg_term self.lossMix = _lossMix @@ -346,13 +564,35 @@ def accuracy(target, output): self.accuracy_fn = accuracy - self.avg_policy_loss = [] - self.avg_value_loss = [] - self.avg_moves_left_loss = [] - self.avg_mse_loss = [] - self.avg_reg_term = [] + # Order must match the order in process_inner_loop + self.train_metrics = [ + Metric('P', 'Policy Loss'), + Metric('V', 'Value Loss'), + Metric('ML', 'Moves Left Loss'), + Metric('Reg', 'Reg term'), + Metric('Total', 'Total Loss'), + Metric( + 'V MSE', 'MSE Loss' + ), # Long name here doesn't mention value for backwards compatibility reasons. + ] self.time_start = None self.last_steps = None + + # Order must match the order in calculate_test_summaries_inner_loop + self.test_metrics = [ + Metric('P', 'Policy Loss'), + Metric('V', 'Value Loss'), + Metric('ML', 'Moves Left Loss'), + Metric( + 'V MSE', 'MSE Loss' + ), # Long name here doesn't mention value for backwards compatibility reasons. + Metric('P Acc', 'Policy Accuracy', suffix='%'), + Metric('V Acc', 'Value Accuracy', suffix='%'), + Metric('ML Mean', 'Moves Left Mean Error'), + Metric('P Entropy', 'Policy Entropy'), + Metric('P UL', 'Policy UL'), + ] + # Set adaptive learning rate during training self.cfg['training']['lr_boundaries'].sort() self.warmup_steps = self.cfg['training'].get('warmup_steps', 0) @@ -388,7 +628,7 @@ def accuracy(target, output): keep_checkpoint_every_n_hours=24, checkpoint_name=self.cfg['name']) - def replace_weights_v2(self, proto_filename, ignore_errors=False): + def replace_weights(self, proto_filename, ignore_errors=False): self.net.parse_proto(proto_filename) filters, blocks = self.net.filters(), self.net.blocks() @@ -471,24 +711,37 @@ def replace_weights_v2(self, proto_filename, ignore_errors=False): # Replace the SWA weights as well, ensuring swa accumulation is reset. if self.swa_enabled: self.swa_count.assign(tf.constant(0.)) - self.update_swa_v2() + self.update_swa() # This should result in identical file to the starting one - # self.save_leelaz_weights_v2('restored.pb.gz') + # self.save_leelaz_weights('restored.pb.gz') - def restore_v2(self): + def restore(self): if self.manager.latest_checkpoint is not None: print("Restoring from {0}".format(self.manager.latest_checkpoint)) self.checkpoint.restore(self.manager.latest_checkpoint) - def process_loop_v2(self, batch_size, test_batches, batch_splits=1): + def process_loop(self, batch_size, test_batches, batch_splits=1): + if self.swa_enabled: + # split half of test_batches between testing regular weights and SWA weights + test_batches //= 2 + # Make sure that ghost batch norm can be applied + if self.virtual_batch_size and batch_size % self.virtual_batch_size != 0: + # Adjust required batch size for batch splitting. + required_factor = self.virtual_batch_size * self.cfg[ + 'training'].get('num_batch_splits', 1) + raise ValueError( + 'batch_size must be a multiple of {}'.format(required_factor)) + # Get the initial steps value in case this is a resume from a step count # which is not a multiple of total_steps. steps = self.global_step.read_value() + self.last_steps = steps + self.time_start = time.time() + self.profiling_start_step = None + total_steps = self.cfg['training']['total_steps'] for _ in range(steps % total_steps, total_steps): - self.process_v2(batch_size, - test_batches, - batch_splits=batch_splits) + self.process(batch_size, test_batches, batch_splits=batch_splits) @tf.function() def read_weights(self): @@ -513,56 +766,69 @@ def process_inner_loop(self, x, y, z, q, m): moves_left_loss = self.moves_left_loss_fn(m, moves_left) else: moves_left_loss = tf.constant(0.) - - total_loss = self.lossMix(policy_loss, value_loss, - moves_left_loss) + reg_term + total_loss = self.lossMix(policy_loss, value_loss, moves_left_loss, + reg_term) if self.loss_scale != 1: total_loss = self.optimizer.get_scaled_loss(total_loss) if self.wdl: mse_loss = self.mse_loss_fn(self.qMix(z, q), value) else: value_loss = self.value_loss_fn(self.qMix(z, q), value) - return policy_loss, value_loss, mse_loss, moves_left_loss, reg_term, tape.gradient( - total_loss, self.model.trainable_weights) - - def process_v2(self, batch_size, test_batches, batch_splits=1): - if not self.time_start: - self.time_start = time.time() + metrics = [ + policy_loss, + value_loss, + moves_left_loss, + reg_term, + total_loss, + # Google's paper scales MSE by 1/4 to a [0, 1] range, so do the same to + # get comparable values. + mse_loss / 4.0, + ] + return metrics, tape.gradient(total_loss, self.model.trainable_weights) - # Get the initial steps value before we do a training step. - steps = self.global_step.read_value() - if not self.last_steps: - self.last_steps = steps + @tf.function() + def strategy_process_inner_loop(self, x, y, z, q, m): + metrics, new_grads = self.strategy.run(self.process_inner_loop, + args=(x, y, z, q, m)) + metrics = [ + self.strategy.reduce(tf.distribute.ReduceOp.MEAN, m, axis=None) + for m in metrics + ] + return metrics, new_grads - if self.swa_enabled: - # split half of test_batches between testing regular weights and SWA weights - test_batches //= 2 + def apply_grads(self, grads, effective_batch_splits): + grads = [ + g[0] + for g in self.aggregator(zip(grads, self.model.trainable_weights)) + ] + if self.loss_scale != 1: + grads = self.optimizer.get_unscaled_gradients(grads) + max_grad_norm = self.cfg['training'].get( + 'max_grad_norm', 10000.0) * effective_batch_splits + grads, grad_norm = tf.clip_by_global_norm(grads, max_grad_norm) + self.optimizer.apply_gradients(zip(grads, + self.model.trainable_weights), + experimental_aggregate_gradients=False) + return grad_norm - # Run test before first step to see delta since end of last run. - if steps % self.cfg['training']['total_steps'] == 0: - # Steps is given as one higher than current in order to avoid it - # being equal to the value the end of a run is stored against. - self.calculate_test_summaries_v2(test_batches, steps + 1) - if self.swa_enabled: - self.calculate_swa_summaries_v2(test_batches, steps + 1) + @tf.function() + def strategy_apply_grads(self, grads, effective_batch_splits): + grad_norm = self.strategy.run(self.apply_grads, + args=(grads, effective_batch_splits)) + grad_norm = self.strategy.reduce(tf.distribute.ReduceOp.MEAN, + grad_norm, + axis=None) + return grad_norm - # Make sure that ghost batch norm can be applied - if self.virtual_batch_size and batch_size % self.virtual_batch_size != 0: - # Adjust required batch size for batch splitting. - required_factor = self.virtual_batch_size * self.cfg[ - 'training'].get('num_batch_splits', 1) - raise ValueError( - 'batch_size must be a multiple of {}'.format(required_factor)) + @tf.function() + def merge_grads(self, grads, new_grads): + return [tf.math.add(a, b) for (a, b) in zip(grads, new_grads)] - # Determine learning rate - lr_values = self.cfg['training']['lr_values'] - lr_boundaries = self.cfg['training']['lr_boundaries'] - steps_total = steps % self.cfg['training']['total_steps'] - self.lr = lr_values[bisect.bisect_right(lr_boundaries, steps_total)] - if self.warmup_steps > 0 and steps < self.warmup_steps: - self.lr = self.lr * tf.cast(steps + 1, - tf.float32) / self.warmup_steps + @tf.function() + def strategy_merge_grads(self, grads, new_grads): + return self.strategy.run(self.merge_grads, args=(grads, new_grads)) + def train_step(self, steps, batch_size, batch_splits): # need to add 1 to steps because steps will be incremented after gradient update if (steps + 1) % self.cfg['training']['train_avg_report_steps'] == 0 or ( @@ -573,32 +839,37 @@ def process_v2(self, batch_size, test_batches, batch_splits=1): grads = None for _ in range(batch_splits): x, y, z, q, m = next(self.train_iter) - policy_loss, value_loss, mse_loss, moves_left_loss, reg_term, new_grads = self.process_inner_loop( - x, y, z, q, m) + if self.strategy is not None: + metrics, new_grads = self.strategy_process_inner_loop( + x, y, z, q, m) + else: + metrics, new_grads = self.process_inner_loop(x, y, z, q, m) if not grads: grads = new_grads else: - grads = [tf.math.add(a, b) for (a, b) in zip(grads, new_grads)] + if self.strategy is not None: + grads = self.strategy_merge_grads(grads, new_grads) + else: + grads = self.merge_grads(grads, new_grads) # Keep running averages - # Google's paper scales MSE by 1/4 to a [0, 1] range, so do the same to - # get comparable values. - mse_loss /= 4.0 - self.avg_policy_loss.append(policy_loss) - if self.wdl: - self.avg_value_loss.append(value_loss) - if self.moves_left: - self.avg_moves_left_loss.append(moves_left_loss) - self.avg_mse_loss.append(mse_loss) - self.avg_reg_term.append(reg_term) + for acc, val in zip(self.train_metrics, metrics): + acc.accumulate(val) # Gradients of batch splits are summed, not averaged like usual, so need to scale lr accordingly to correct for this. - self.active_lr = self.lr / batch_splits - if self.loss_scale != 1: - grads = self.optimizer.get_unscaled_gradients(grads) - max_grad_norm = self.cfg['training'].get('max_grad_norm', - 10000.0) * batch_splits - grads, grad_norm = tf.clip_by_global_norm(grads, max_grad_norm) - self.optimizer.apply_gradients(zip(grads, - self.model.trainable_weights)) + effective_batch_splits = batch_splits + if self.strategy is not None: + effective_batch_splits = batch_splits * self.strategy.num_replicas_in_sync + self.active_lr.assign(self.lr / effective_batch_splits) + if self.update_lr_manually: + self.orig_optimizer.learning_rate = self.active_lr + if self.strategy is not None: + grad_norm = self.strategy_apply_grads(grads, + effective_batch_splits) + else: + grad_norm = self.apply_grads(grads, effective_batch_splits) + + # Note: grads variable at this point has not been unscaled or + # had clipping applied. Since no code after this point depends + # upon that it seems fine for now. # Update steps. self.global_step.assign_add(1) @@ -607,9 +878,6 @@ def process_v2(self, batch_size, test_batches, batch_splits=1): if steps % self.cfg['training'][ 'train_avg_report_steps'] == 0 or steps % self.cfg['training'][ 'total_steps'] == 0: - pol_loss_w = self.cfg['training']['policy_loss_weight'] - val_loss_w = self.cfg['training']['value_loss_weight'] - moves_loss_w = self.cfg['training']['moves_left_loss_weight'] time_end = time.time() speed = 0 if self.time_start: @@ -617,63 +885,86 @@ def process_v2(self, batch_size, test_batches, batch_splits=1): steps_elapsed = steps - self.last_steps speed = batch_size * (tf.cast(steps_elapsed, tf.float32) / elapsed) - avg_policy_loss = np.mean(self.avg_policy_loss or [0]) - avg_moves_left_loss = np.mean(self.avg_moves_left_loss or [0]) - avg_value_loss = np.mean(self.avg_value_loss or [0]) - avg_mse_loss = np.mean(self.avg_mse_loss or [0]) - avg_reg_term = np.mean(self.avg_reg_term or [0]) - print( - "step {}, lr={:g} policy={:g} value={:g} mse={:g} moves={:g} reg={:g} total={:g} ({:g} pos/s)" - .format( - steps, self.lr, avg_policy_loss, avg_value_loss, - avg_mse_loss, avg_moves_left_loss, avg_reg_term, - pol_loss_w * avg_policy_loss + - val_loss_w * avg_value_loss + avg_reg_term + - moves_loss_w * avg_moves_left_loss, speed)) + print("step {}, lr={:g}".format(steps, self.lr), end='') + for metric in self.train_metrics: + print(" {}={:g}{}".format(metric.short_name, metric.get(), + metric.suffix), + end='') + print(" ({:g} pos/s)".format(speed)) after_weights = self.read_weights() with self.train_writer.as_default(): - tf.summary.scalar("Policy Loss", avg_policy_loss, step=steps) - tf.summary.scalar("Value Loss", avg_value_loss, step=steps) - if self.moves_left: - tf.summary.scalar("Moves Left Loss", - avg_moves_left_loss, + for metric in self.train_metrics: + tf.summary.scalar(metric.long_name, + metric.get(), step=steps) - tf.summary.scalar("Reg term", avg_reg_term, step=steps) tf.summary.scalar("LR", self.lr, step=steps) tf.summary.scalar("Gradient norm", - grad_norm / batch_splits, + grad_norm / effective_batch_splits, step=steps) - tf.summary.scalar("MSE Loss", avg_mse_loss, step=steps) - self.compute_update_ratio_v2(before_weights, after_weights, - steps) + self.compute_update_ratio(before_weights, after_weights, steps) self.train_writer.flush() + self.time_start = time_end self.last_steps = steps - self.avg_policy_loss = [] - self.avg_moves_left_loss = [] - self.avg_value_loss = [] - self.avg_mse_loss = [] - self.avg_reg_term = [] + for metric in self.train_metrics: + metric.reset() + return steps + + def process(self, batch_size, test_batches, batch_splits): + # Get the initial steps value before we do a training step. + steps = self.global_step.read_value() + + # By default disabled since 0 != 10. + if steps % self.cfg['training'].get('profile_step_freq', + 1) == self.cfg['training'].get( + 'profile_step_offset', 10): + self.profiling_start_step = steps + tf.profiler.experimental.start( + os.path.join(os.getcwd(), + "leelalogs/{}-profile".format(self.cfg['name']))) + + # Run test before first step to see delta since end of last run. + if steps % self.cfg['training']['total_steps'] == 0: + with tf.profiler.experimental.Trace("Test", step_num=steps + 1): + # Steps is given as one higher than current in order to avoid it + # being equal to the value the end of a run is stored against. + self.calculate_test_summaries(test_batches, steps + 1) + if self.swa_enabled: + self.calculate_swa_summaries(test_batches, steps + 1) + + # Determine learning rate + lr_values = self.cfg['training']['lr_values'] + lr_boundaries = self.cfg['training']['lr_boundaries'] + steps_total = steps % self.cfg['training']['total_steps'] + self.lr = lr_values[bisect.bisect_right(lr_boundaries, steps_total)] + if self.warmup_steps > 0 and steps < self.warmup_steps: + self.lr = self.lr * tf.cast(steps + 1, + tf.float32) / self.warmup_steps + + with tf.profiler.experimental.Trace("Train", step_num=steps): + steps = self.train_step(steps, batch_size, batch_splits) if self.swa_enabled and steps % self.cfg['training']['swa_steps'] == 0: - self.update_swa_v2() + self.update_swa() # Calculate test values every 'test_steps', but also ensure there is - # one at the final step so the delta to the first step can be calculted. + # one at the final step so the delta to the first step can be calculated. if steps % self.cfg['training']['test_steps'] == 0 or steps % self.cfg[ 'training']['total_steps'] == 0: - self.calculate_test_summaries_v2(test_batches, steps) - if self.swa_enabled: - self.calculate_swa_summaries_v2(test_batches, steps) + with tf.profiler.experimental.Trace("Test", step_num=steps): + self.calculate_test_summaries(test_batches, steps) + if self.swa_enabled: + self.calculate_swa_summaries(test_batches, steps) if self.validation_dataset is not None and ( steps % self.cfg['training']['validation_steps'] == 0 or steps % self.cfg['training']['total_steps'] == 0): - if self.swa_enabled: - self.calculate_swa_validations_v2(steps) - else: - self.calculate_test_validations_v2(steps) + with tf.profiler.experimental.Trace("Validate", step_num=steps): + if self.swa_enabled: + self.calculate_swa_validations(steps) + else: + self.calculate_test_validations(steps) # Save session and weights at end, and also optionally every 'checkpoint_steps'. if steps % self.cfg['training']['total_steps'] == 0 or ( @@ -687,17 +978,24 @@ def process_v2(self, batch_size, test_batches, batch_splits=1): leela_path = path + "-" + str(evaled_steps) swa_path = path + "-swa-" + str(evaled_steps) self.net.pb.training_params.training_steps = evaled_steps - self.save_leelaz_weights_v2(leela_path) + self.save_leelaz_weights(leela_path) if self.swa_enabled: - self.save_swa_weights_v2(swa_path) + self.save_swa_weights(swa_path) + + if self.profiling_start_step is not None and ( + steps >= self.profiling_start_step + + self.cfg['training'].get('profile_step_count', 0) + or steps % self.cfg['training']['total_steps'] == 0): + tf.profiler.experimental.stop() + self.profiling_start_step = None - def calculate_swa_summaries_v2(self, test_batches, steps): + def calculate_swa_summaries(self, test_batches, steps): backup = self.read_weights() for (swa, w) in zip(self.swa_weights, self.model.weights): w.assign(swa.read_value()) true_test_writer, self.test_writer = self.test_writer, self.swa_writer print('swa', end=' ') - self.calculate_test_summaries_v2(test_batches, steps) + self.calculate_test_summaries(test_batches, steps) self.test_writer = true_test_writer for (old, w) in zip(backup, self.model.weights): w.assign(old) @@ -726,170 +1024,98 @@ def calculate_test_summaries_inner_loop(self, x, y, z, q, m): else: moves_left_loss = tf.constant(0.) moves_left_mean_error = tf.constant(0.) + metrics = [ + policy_loss, + value_loss, + moves_left_loss, + mse_loss / 4, + policy_accuracy * 100, + value_accuracy * 100, + moves_left_mean_error, + policy_entropy, + policy_ul, + ] + return metrics - return policy_loss, value_loss, moves_left_loss, mse_loss, policy_accuracy, value_accuracy, moves_left_mean_error, policy_entropy, policy_ul - - def calculate_test_summaries_v2(self, test_batches, steps): - sum_policy_accuracy = 0 - sum_value_accuracy = 0 - sum_moves_left = 0 - sum_moves_left_mean_error = 0 - sum_mse = 0 - sum_policy = 0 - sum_value = 0 - sum_policy_entropy = 0 - sum_policy_ul = 0 + @tf.function() + def strategy_calculate_test_summaries_inner_loop(self, x, y, z, q, m): + metrics = self.strategy.run(self.calculate_test_summaries_inner_loop, + args=(x, y, z, q, m)) + metrics = [ + self.strategy.reduce(tf.distribute.ReduceOp.MEAN, m, axis=None) + for m in metrics + ] + return metrics + + def calculate_test_summaries(self, test_batches, steps): + for metric in self.test_metrics: + metric.reset() for _ in range(0, test_batches): x, y, z, q, m = next(self.test_iter) - policy_loss, value_loss, moves_left_loss, mse_loss, policy_accuracy, value_accuracy, moves_left_mean_error, policy_entropy, policy_ul = self.calculate_test_summaries_inner_loop( - x, y, z, q, m) - sum_policy_accuracy += policy_accuracy - sum_policy_entropy += policy_entropy - sum_policy_ul += policy_ul - sum_mse += mse_loss - sum_policy += policy_loss - if self.wdl: - sum_value_accuracy += value_accuracy - sum_value += value_loss - if self.moves_left: - sum_moves_left += moves_left_loss - sum_moves_left_mean_error += moves_left_mean_error - sum_policy_accuracy /= test_batches - sum_policy_accuracy *= 100 - sum_policy /= test_batches - sum_policy_entropy /= test_batches - sum_policy_ul /= test_batches - sum_value /= test_batches - if self.wdl: - sum_value_accuracy /= test_batches - sum_value_accuracy *= 100 - # Additionally rescale to [0, 1] so divide by 4 - sum_mse /= (4.0 * test_batches) - if self.moves_left: - sum_moves_left /= test_batches - sum_moves_left_mean_error /= test_batches + if self.strategy is not None: + metrics = self.strategy_calculate_test_summaries_inner_loop( + x, y, z, q, m) + else: + metrics = self.calculate_test_summaries_inner_loop( + x, y, z, q, m) + for acc, val in zip(self.test_metrics, metrics): + acc.accumulate(val) self.net.pb.training_params.learning_rate = self.lr - self.net.pb.training_params.mse_loss = sum_mse - self.net.pb.training_params.policy_loss = sum_policy + self.net.pb.training_params.mse_loss = self.test_metrics[3].get() + self.net.pb.training_params.policy_loss = self.test_metrics[0].get() # TODO store value and value accuracy in pb - self.net.pb.training_params.accuracy = sum_policy_accuracy + self.net.pb.training_params.accuracy = self.test_metrics[4].get() with self.test_writer.as_default(): - tf.summary.scalar("Policy Loss", sum_policy, step=steps) - tf.summary.scalar("Value Loss", sum_value, step=steps) - tf.summary.scalar("MSE Loss", sum_mse, step=steps) - tf.summary.scalar("Policy Accuracy", - sum_policy_accuracy, - step=steps) - tf.summary.scalar("Policy Entropy", sum_policy_entropy, step=steps) - tf.summary.scalar("Policy UL", sum_policy_ul, step=steps) - if self.wdl: - tf.summary.scalar("Value Accuracy", - sum_value_accuracy, - step=steps) - if self.moves_left: - tf.summary.scalar("Moves Left Loss", - sum_moves_left, - step=steps) - tf.summary.scalar("Moves Left Mean Error", - sum_moves_left_mean_error, - step=steps) + for metric in self.test_metrics: + tf.summary.scalar(metric.long_name, metric.get(), step=steps) for w in self.model.weights: tf.summary.histogram(w.name, w, step=steps) self.test_writer.flush() - print("step {}, policy={:g} value={:g} policy accuracy={:g}% value accuracy={:g}% mse={:g} policy entropy={:g} policy ul={:g}".\ - format(steps, sum_policy, sum_value, sum_policy_accuracy, sum_value_accuracy, sum_mse, sum_policy_entropy, sum_policy_ul), end = '') - - if self.moves_left: - print(" moves={:g} moves mean={:g}".format( - sum_moves_left, sum_moves_left_mean_error)) - else: - print() + print("step {},".format(steps), end='') + for metric in self.test_metrics: + print(" {}={:g}{}".format(metric.short_name, metric.get(), + metric.suffix), + end='') + print() - def calculate_swa_validations_v2(self, steps): + def calculate_swa_validations(self, steps): backup = self.read_weights() for (swa, w) in zip(self.swa_weights, self.model.weights): w.assign(swa.read_value()) true_validation_writer, self.validation_writer = self.validation_writer, self.swa_validation_writer print('swa', end=' ') - self.calculate_test_validations_v2(steps) + self.calculate_test_validations(steps) self.validation_writer = true_validation_writer for (old, w) in zip(backup, self.model.weights): w.assign(old) - def calculate_test_validations_v2(self, steps): - sum_policy_accuracy = 0 - sum_value_accuracy = 0 - sum_moves_left = 0 - sum_moves_left_mean_error = 0 - sum_mse = 0 - sum_policy = 0 - sum_value = 0 - sum_policy_entropy = 0 - sum_policy_ul = 0 - counter = 0 + def calculate_test_validations(self, steps): + for metric in self.test_metrics: + metric.reset() for (x, y, z, q, m) in self.validation_dataset: - policy_loss, value_loss, moves_left_loss, mse_loss, policy_accuracy, value_accuracy, moves_left_mean_error, policy_entropy, policy_ul = self.calculate_test_summaries_inner_loop( - x, y, z, q, m) - sum_policy_accuracy += policy_accuracy - sum_policy_entropy += policy_entropy - sum_policy_ul += policy_ul - sum_mse += mse_loss - sum_policy += policy_loss - if self.moves_left: - sum_moves_left += moves_left_loss - sum_moves_left_mean_error += moves_left_mean_error - counter += 1 - if self.wdl: - sum_value_accuracy += value_accuracy - sum_value += value_loss - sum_policy_accuracy /= counter - sum_policy_accuracy *= 100 - sum_policy /= counter - sum_policy_entropy /= counter - sum_policy_ul /= counter - sum_value /= counter - if self.wdl: - sum_value_accuracy /= counter - sum_value_accuracy *= 100 - if self.moves_left: - sum_moves_left /= counter - sum_moves_left_mean_error /= counter - # Additionally rescale to [0, 1] so divide by 4 - sum_mse /= (4.0 * counter) + if self.strategy is not None: + metrics = self.strategy_calculate_test_summaries_inner_loop( + x, y, z, q, m) + else: + metrics = self.calculate_test_summaries_inner_loop( + x, y, z, q, m) + for acc, val in zip(self.test_metrics, metrics): + acc.accumulate(val) with self.validation_writer.as_default(): - tf.summary.scalar("Policy Loss", sum_policy, step=steps) - tf.summary.scalar("Value Loss", sum_value, step=steps) - tf.summary.scalar("MSE Loss", sum_mse, step=steps) - tf.summary.scalar("Policy Accuracy", - sum_policy_accuracy, - step=steps) - tf.summary.scalar("Policy Entropy", sum_policy_entropy, step=steps) - tf.summary.scalar("Policy UL", sum_policy_ul, step=steps) - if self.wdl: - tf.summary.scalar("Value Accuracy", - sum_value_accuracy, - step=steps) - if self.moves_left: - tf.summary.scalar("Moves Left Loss", - sum_moves_left, - step=steps) - tf.summary.scalar("Moves Left Mean Error", - sum_moves_left_mean_error, - step=steps) + for metric in self.test_metrics: + tf.summary.scalar(metric.long_name, metric.get(), step=steps) self.validation_writer.flush() - print("step {}, validation: policy={:g} value={:g} policy accuracy={:g}% value accuracy={:g}% mse={:g} policy entropy={:g} policy ul={:g}".\ - format(steps, sum_policy, sum_value, sum_policy_accuracy, sum_value_accuracy, sum_mse, sum_policy_entropy, sum_policy_ul), end='') - - if self.moves_left: - print(" moves={:g} moves mean={:g}".format( - sum_moves_left, sum_moves_left_mean_error)) - else: - print() + print("step {}, validation:".format(steps), end='') + for metric in self.test_metrics: + print(" {}={:g}{}".format(metric.short_name, metric.get(), + metric.suffix), + end='') + print() @tf.function() - def compute_update_ratio_v2(self, before_weights, after_weights, steps): + def compute_update_ratio(self, before_weights, after_weights, steps): """Compute the ratio of gradient norm to weight norm. Adapted from https://github.com/tensorflow/minigo/blob/c923cd5b11f7d417c9541ad61414bf175a84dc31/dual_net.py#L567 @@ -918,29 +1144,29 @@ def compute_update_ratio_v2(self, before_weights, after_weights, steps): buckets=1000, step=steps) - def update_swa_v2(self): + def update_swa(self): num = self.swa_count.read_value() for (w, swa) in zip(self.model.weights, self.swa_weights): swa.assign(swa.read_value() * (num / (num + 1.)) + w.read_value() * (1. / (num + 1.))) self.swa_count.assign(min(num + 1., self.swa_max_n)) - def save_swa_weights_v2(self, filename): + def save_swa_weights(self, filename): backup = self.read_weights() for (swa, w) in zip(self.swa_weights, self.model.weights): w.assign(swa.read_value()) - self.save_leelaz_weights_v2(filename) + self.save_leelaz_weights(filename) for (old, w) in zip(backup, self.model.weights): w.assign(old) - def save_leelaz_weights_v2(self, filename): + def save_leelaz_weights(self, filename): numpy_weights = [] for weight in self.model.weights: numpy_weights.append([weight.name, weight.numpy()]) self.net.fill_net_v2(numpy_weights) self.net.save_proto(filename) - def batch_norm_v2(self, input, name, scale=False): + def batch_norm(self, input, name, scale=False): if self.renorm_enabled: clipping = { "rmin": 1.0 / self.renorm_max_r, @@ -966,28 +1192,28 @@ def batch_norm_v2(self, input, name, scale=False): virtual_batch_size=self.virtual_batch_size, name=name)(input) - def squeeze_excitation_v2(self, inputs, channels, name): + def squeeze_excitation(self, inputs, channels, name): assert channels % self.SE_ratio == 0 pooled = tf.keras.layers.GlobalAveragePooling2D( data_format='channels_first')(inputs) - squeezed = tf.keras.layers.Activation('relu')(tf.keras.layers.Dense( - channels // self.SE_ratio, - kernel_initializer='glorot_normal', - kernel_regularizer=self.l2reg, - name=name + '/se/dense1')(pooled)) + squeezed = tf.keras.layers.Activation(self.DEFAULT_ACTIVATION)( + tf.keras.layers.Dense(channels // self.SE_ratio, + kernel_initializer='glorot_normal', + kernel_regularizer=self.l2reg, + name=name + '/se/dense1')(pooled)) excited = tf.keras.layers.Dense(2 * channels, kernel_initializer='glorot_normal', kernel_regularizer=self.l2reg, name=name + '/se/dense2')(squeezed) return ApplySqueezeExcitation()([inputs, excited]) - def conv_block_v2(self, - inputs, - filter_size, - output_channels, - name, - bn_scale=False): + def conv_block(self, + inputs, + filter_size, + output_channels, + name, + bn_scale=False): conv = tf.keras.layers.Conv2D(output_channels, filter_size, use_bias=False, @@ -996,10 +1222,10 @@ def conv_block_v2(self, kernel_regularizer=self.l2reg, data_format='channels_first', name=name + '/conv2d')(inputs) - return tf.keras.layers.Activation('relu')(self.batch_norm_v2( - conv, name=name + '/bn', scale=bn_scale)) + return tf.keras.layers.Activation(self.DEFAULT_ACTIVATION)( + self.batch_norm(conv, name=name + '/bn', scale=bn_scale)) - def residual_block_v2(self, inputs, channels, name): + def residual_block(self, inputs, channels, name): conv1 = tf.keras.layers.Conv2D(channels, 3, use_bias=False, @@ -1008,8 +1234,8 @@ def residual_block_v2(self, inputs, channels, name): kernel_regularizer=self.l2reg, data_format='channels_first', name=name + '/1/conv2d')(inputs) - out1 = tf.keras.layers.Activation('relu')(self.batch_norm_v2( - conv1, name + '/1/bn', scale=False)) + out1 = tf.keras.layers.Activation(self.DEFAULT_ACTIVATION)( + self.batch_norm(conv1, name + '/1/bn', scale=False)) conv2 = tf.keras.layers.Conv2D(channels, 3, use_bias=False, @@ -1018,32 +1244,296 @@ def residual_block_v2(self, inputs, channels, name): kernel_regularizer=self.l2reg, data_format='channels_first', name=name + '/2/conv2d')(out1) - out2 = self.squeeze_excitation_v2(self.batch_norm_v2(conv2, - name + '/2/bn', - scale=True), - channels, - name=name + '/se') - return tf.keras.layers.Activation('relu')(tf.keras.layers.add( - [inputs, out2])) - - def construct_net_v2(self, inputs): - flow = self.conv_block_v2(inputs, - filter_size=3, - output_channels=self.RESIDUAL_FILTERS, - name='input', - bn_scale=True) + + out2 = self.squeeze_excitation(self.batch_norm(conv2, + name + '/2/bn', + scale=True), + channels, + name=name + '/se') + return tf.keras.layers.Activation(self.DEFAULT_ACTIVATION)( + tf.keras.layers.add([inputs, out2])) + + @staticmethod + def split_heads(inputs, batch_size: int, num_heads: int, depth: int): + if num_heads < 2: + return inputs + reshaped = tf.reshape(inputs, (batch_size, 64, num_heads, depth)) + # (batch_size, num_heads, 64, depth) + return tf.transpose(reshaped, perm=[0, 2, 1, 3]) + + def scaled_dot_product_attention(self, + q, + k, + v, + name: str = None, + inputs=None): + + # 0 h 64 d, 0 h d 64 + matmul_qk = tf.matmul(q, k, transpose_b=True) + dk = tf.cast(tf.shape(k)[-1], self.model_dtype) + scaled_attention_logits = matmul_qk / tf.math.sqrt(dk) + heads = scaled_attention_logits.shape[1] + + if self.use_smolgen: + smolgen_weights = self.smolgen_weights( + inputs, + heads, + self.smolgen_hidden_channels, + self.smolgen_hidden_sz, + self.smolgen_gen_sz, + name=name + '/smolgen', + activation=self.smolgen_activation) + scaled_attention_logits = scaled_attention_logits + smolgen_weights + + attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1) + output = tf.matmul(attention_weights, v) + return output, scaled_attention_logits + + # multi-head attention in encoder blocks + + def mha(self, inputs, emb_size: int, d_model: int, num_heads: int, + initializer, name: str): + assert d_model % num_heads == 0 + depth = d_model // num_heads + # query, key, and value vectors for self-attention + # inputs b, 64, sz + q = tf.keras.layers.Dense(d_model, + name=name + '/wq', + kernel_initializer='glorot_normal')(inputs) + k = tf.keras.layers.Dense(d_model, + name=name + '/wk', + kernel_initializer='glorot_normal')(inputs) + v = tf.keras.layers.Dense(d_model, + name=name + '/wv', + kernel_initializer=initializer)(inputs) + + # split q, k and v into smaller vectors of size 'depth' -- one for each head in multi-head attention + batch_size = tf.shape(q)[0] + q = self.split_heads(q, batch_size, num_heads, depth) + k = self.split_heads(k, batch_size, num_heads, depth) + v = self.split_heads(v, batch_size, num_heads, depth) + + scaled_attention, attention_weights = self.scaled_dot_product_attention( + q, k, v, name=name, inputs=inputs) + if num_heads > 1: + scaled_attention = tf.transpose(scaled_attention, + perm=[0, 2, 1, 3]) + scaled_attention = tf.reshape( + scaled_attention, + (batch_size, -1, d_model)) # concatenate heads + + # final dense layer + output = tf.keras.layers.Dense( + emb_size, name=name + "/dense", + kernel_initializer=initializer)(scaled_attention) + return output, attention_weights + + # 2-layer dense feed-forward network in encoder blocks + def ffn(self, inputs, emb_size: int, dff: int, initializer, name: str): + if self.encoder_layers > 0: + activation = square_relu if self.square_relu_ffn else tf.keras.activations.get( + self.DEFAULT_ACTIVATION) + else: + activation = "selu" + dense1 = tf.keras.layers.Dense(dff, + name=name + "/dense1", + kernel_initializer=initializer, + activation=activation)(inputs) + out = tf.keras.layers.Dense(emb_size, + name=name + "/dense2", + kernel_initializer=initializer)(dense1) + return out + + def encoder_layer(self, inputs, emb_size: int, d_model: int, + num_heads: int, dff: int, name: str): + initializer = None + if self.encoder_layers > 0: + # DeepNorm + alpha = tf.cast(tf.math.pow(2. * self.encoder_layers, 0.25), + self.model_dtype) + beta = tf.cast(tf.math.pow(8. * self.encoder_layers, -0.25), + self.model_dtype) + xavier_norm = tf.keras.initializers.VarianceScaling( + scale=beta, mode='fan_avg', distribution='truncated_normal') + initializer = xavier_norm + else: + alpha = 1 + initializer = "glorot_normal" + # multihead attention + attn_output, attn_wts = self.mha(inputs, + emb_size, + d_model, + num_heads, + initializer, + name=name + "/mha") + # dropout for weight regularization + attn_output = tf.keras.layers.Dropout(self.dropout_rate, + name=name + + "/dropout1")(attn_output) + # skip connection + layernorm + out1 = tf.keras.layers.LayerNormalization( + epsilon=1e-6, name=name + "/ln1")(inputs * alpha + attn_output) + # feed-forward network + ffn_output = self.ffn(out1, + emb_size, + dff, + initializer, + name=name + "/ffn") + ffn_output = tf.keras.layers.Dropout(self.dropout_rate, + name=name + + "/dropout2")(ffn_output) + out2 = tf.keras.layers.LayerNormalization( + epsilon=1e-6, name=name + "/ln2")(out1 * alpha + ffn_output) + return out2, attn_wts + + def smolgen_weights(self, + inputs, + heads: int, + hidden_channels: int, + hidden_sz: int, + gen_sz: int, + name: str, + activation='swish'): + compressed = tf.keras.layers.Dense(hidden_channels, + name=name + '/compress', + use_bias=False)(inputs) + compressed = tf.reshape(compressed, [-1, 64 * hidden_channels]) + + hidden = tf.keras.layers.Dense(hidden_sz, + name=name + '/hidden1_dense', + activation=activation)(compressed) + hidden = tf.keras.layers.LayerNormalization(name=name + + '/hidden1_ln')(hidden) + + gen_from = tf.keras.layers.Dense(heads * gen_sz, + name=name + '/gen_from', + activation=activation)(hidden) + gen_from = tf.keras.layers.LayerNormalization(name=name + + '/gen_from_ln')(gen_from) + gen_from = tf.reshape(gen_from, [-1, heads, gen_sz]) + out = self.smol_weight_gen_dense(gen_from) + return tf.reshape(out, [-1, heads, 64, 64]) + + def create_residual_body(self, inputs): + flow = self.conv_block(inputs, + filter_size=3, + output_channels=self.RESIDUAL_FILTERS, + name='input', + bn_scale=True) for i in range(self.RESIDUAL_BLOCKS): - flow = self.residual_block_v2(flow, - self.RESIDUAL_FILTERS, - name='residual_{}'.format(i + 1)) + flow = self.residual_block(flow, + self.RESIDUAL_FILTERS, + name='residual_{}'.format(i + 1)) + return flow + + def create_encoder_body(self, inputs, embedding_size): + # Policy head + assert self.POLICY_HEAD == pb.NetworkFormat.POLICY_ATTENTION + + # do some input processing + if self.use_smolgen: + self.smol_weight_gen_dense = tf.keras.layers.Dense( + 64 * 64, name='smol_weight_gen', use_bias=False) + + flow = tf.transpose(inputs, perm=[0, 2, 3, 1]) + flow = tf.reshape(flow, [-1, 64, tf.shape(inputs)[1]]) + # add positional encoding for each square to the input + if self.arc_encoding: + self.POS_ENC = apm.make_pos_enc() + positional_encoding = tf.broadcast_to( + tf.convert_to_tensor(self.POS_ENC, dtype=flow.dtype), + [tf.shape(flow)[0], 64, + tf.shape(self.POS_ENC)[2]]) + flow = tf.concat([flow, positional_encoding], axis=2) + + # square embedding + flow = tf.keras.layers.Dense(embedding_size, + kernel_initializer='glorot_normal', + kernel_regularizer=self.l2reg, + activation=self.DEFAULT_ACTIVATION, + name='embedding')(flow) + + # !!! input gate + flow = ma_gating(flow, name='embedding') + attn_wts = [] + for i in range(self.encoder_layers): + flow, attn_wts_l = self.encoder_layer(flow, + embedding_size, + self.encoder_d_model, + self.encoder_heads, + self.encoder_dff, + name='encoder_{}'.format(i + + 1)) + attn_wts.append(attn_wts_l) + return flow, attn_wts + + def apply_promotion_logits(self, queries, keys, attn_wts): + # PAWN PROMOTION: create promotion logits using scalar offsets generated from the promotion-rank keys + dk = tf.math.sqrt(tf.cast(tf.shape(keys)[-1], + self.model_dtype)) # constant for scaling + promotion_keys = keys[:, -8:, :] + # queen, rook, bishop, knight order + promotion_offsets = tf.keras.layers.Dense( + 4, + kernel_initializer='glorot_normal', + name='policy/attention/ppo', + use_bias=False)(promotion_keys) + promotion_offsets = tf.transpose(promotion_offsets, + perm=[0, 2, 1]) * dk # Bx4x8 + # knight offset is added to the other three + promotion_offsets = promotion_offsets[:, : + 3, :] + promotion_offsets[:, + 3:4, :] + + # POLICY SELF-ATTENTION: self-attention weights are interpreted as from->to policy + matmul_qk = tf.matmul( + queries, keys, + transpose_b=True) # Bx64x64 (from 64 queries, 64 keys) + + # q, r, and b promotions are offset from the default promotion logit (knight) + n_promo_logits = matmul_qk[:, -16:-8, + -8:] # default traversals from penultimate rank to promotion rank + q_promo_logits = tf.expand_dims(n_promo_logits + + promotion_offsets[:, 0:1, :], + axis=3) # Bx8x8x1 + r_promo_logits = tf.expand_dims(n_promo_logits + + promotion_offsets[:, 1:2, :], + axis=3) + b_promo_logits = tf.expand_dims(n_promo_logits + + promotion_offsets[:, 2:3, :], + axis=3) + promotion_logits = tf.concat( + [q_promo_logits, r_promo_logits, b_promo_logits], + axis=3) # Bx8x8x3 + promotion_logits = tf.reshape( + promotion_logits, + [-1, 8, 24]) # logits now alternate a7a8q,a7a8r,a7a8b,..., + + # scale the logits by dividing them by sqrt(d_model) to stabilize gradients + promotion_logits = promotion_logits / dk # Bx8x24 (8 from-squares, 3x8 promotions) + policy_attn_logits = matmul_qk / dk # Bx64x64 (64 from-squares, 64 to-squares) + + attn_wts.append(promotion_logits) + attn_wts.append(policy_attn_logits) + + # APPLY POLICY MAP: output becomes Bx1856 + h_fc1 = ApplyAttentionPolicyMap()(policy_attn_logits, promotion_logits) + return h_fc1 + + def construct_net(self, inputs, name=''): + + if self.encoder_layers > 0: + flow, attn_wts = self.create_encoder_body(inputs, + self.embedding_size) + else: + flow = self.create_residual_body(inputs) # Policy head if self.POLICY_HEAD == pb.NetworkFormat.POLICY_CONVOLUTION: - conv_pol = self.conv_block_v2( - flow, - filter_size=3, - output_channels=self.RESIDUAL_FILTERS, - name='policy1') + conv_pol = self.conv_block(flow, + filter_size=3, + output_channels=self.RESIDUAL_FILTERS, + name='policy1') conv_pol2 = tf.keras.layers.Conv2D( 80, 3, @@ -1056,30 +1546,78 @@ def construct_net_v2(self, inputs): name='policy')(conv_pol) h_fc1 = ApplyPolicyMap()(conv_pol2) elif self.POLICY_HEAD == pb.NetworkFormat.POLICY_CLASSICAL: - conv_pol = self.conv_block_v2(flow, - filter_size=1, - output_channels=self.policy_channels, - name='policy') + conv_pol = self.conv_block(flow, + filter_size=1, + output_channels=self.policy_channels, + name='policy') h_conv_pol_flat = tf.keras.layers.Flatten()(conv_pol) h_fc1 = tf.keras.layers.Dense(1858, kernel_initializer='glorot_normal', kernel_regularizer=self.l2reg, bias_regularizer=self.l2reg, name='policy/dense')(h_conv_pol_flat) + elif self.POLICY_HEAD == pb.NetworkFormat.POLICY_ATTENTION: + if self.encoder_layers == 0: + attn_wts = [] + if self.RESIDUAL_BLOCKS > 0: + # transpose and reshape + tokens = tf.transpose(flow, perm=[0, 2, 3, 1]) + tokens = tf.reshape(tokens, [-1, 64, self.RESIDUAL_FILTERS]) + embed_activation = 'selu' + else: + tokens = flow + embed_activation = self.DEFAULT_ACTIVATION + # SQUARE EMBEDDING: found to increase attention head performance + tokens = tf.keras.layers.Dense(self.pol_embedding_size, + kernel_initializer='glorot_normal', + kernel_regularizer=self.l2reg, + activation=embed_activation, + name='policy/embedding')(tokens) + if self.RESIDUAL_BLOCKS > 0: + # ENCODER LAYERS: intermediate layers of self-attention with residual connections + for i in range(self.pol_encoder_layers): + tokens, attn_wts_l = self.encoder_layer( + tokens, + self.pol_embedding_size, + self.pol_encoder_d_model, + self.pol_encoder_heads, + self.pol_encoder_dff, + name='policy/enc_layer_{}'.format(i + 1)) + attn_wts.append(attn_wts_l) + + # create queries and keys for policy self-attention + queries = tf.keras.layers.Dense(self.policy_d_model, + kernel_initializer='glorot_normal', + name='policy/attention/wq')(tokens) + keys = tf.keras.layers.Dense(self.policy_d_model, + kernel_initializer='glorot_normal', + name='policy/attention/wk')(tokens) + + h_fc1 = self.apply_promotion_logits(queries, keys, attn_wts) + else: raise ValueError("Unknown policy head type {}".format( self.POLICY_HEAD)) # Value head - conv_val = self.conv_block_v2(flow, - filter_size=1, - output_channels=32, - name='value') + if self.encoder_layers > 0: + conv_val = tf.keras.layers.Dense( + self.val_embedding_size, + kernel_initializer='glorot_normal', + kernel_regularizer=self.l2reg, + activation=self.DEFAULT_ACTIVATION, + name='value/embedding')(flow) + else: + conv_val = self.conv_block(flow, + filter_size=1, + output_channels=32, + name='value') + h_conv_val_flat = tf.keras.layers.Flatten()(conv_val) h_fc2 = tf.keras.layers.Dense(128, kernel_initializer='glorot_normal', kernel_regularizer=self.l2reg, - activation='relu', + activation=self.DEFAULT_ACTIVATION, name='value/dense1')(h_conv_val_flat) if self.wdl: h_fc3 = tf.keras.layers.Dense(3, @@ -1096,16 +1634,24 @@ def construct_net_v2(self, inputs): # Moves left head if self.moves_left: - conv_mov = self.conv_block_v2(flow, - filter_size=1, - output_channels=8, - name='moves_left') + if self.encoder_layers > 0: + conv_mov = tf.keras.layers.Dense( + self.mov_embedding_size, + kernel_initializer='glorot_normal', + kernel_regularizer=self.l2reg, + activation=self.DEFAULT_ACTIVATION, + name='moves_left/embedding')(flow) + else: + conv_mov = self.conv_block(flow, + filter_size=1, + output_channels=8, + name='moves_left') h_conv_mov_flat = tf.keras.layers.Flatten()(conv_mov) h_fc4 = tf.keras.layers.Dense( 128, kernel_initializer='glorot_normal', kernel_regularizer=self.l2reg, - activation='relu', + activation=self.DEFAULT_ACTIVATION, name='moves_left/dense1')(h_conv_mov_flat) h_fc5 = tf.keras.layers.Dense(1, @@ -1116,4 +1662,15 @@ def construct_net_v2(self, inputs): else: h_fc5 = None - return h_fc1, h_fc3, h_fc5 + # attention weights added as optional output for analysis -- ignored by backend + if self.POLICY_HEAD == pb.NetworkFormat.POLICY_ATTENTION: + if self.moves_left: + outputs = [h_fc1, h_fc3, h_fc5, attn_wts] + else: + outputs = [h_fc1, h_fc3, attn_wts] + elif self.moves_left: + outputs = [h_fc1, h_fc3, h_fc5] + else: + outputs = [h_fc1, h_fc3] + + return outputs diff --git a/tf/train.py b/tf/train.py index a4a1745d..2935c313 100644 --- a/tf/train.py +++ b/tf/train.py @@ -24,12 +24,9 @@ import gzip import random import multiprocessing as mp -import tensorflow as tf -from tfprocess import TFProcess from chunkparser import ChunkParser SKIP = 32 -SKIP_MULTIPLE = 1024 def get_chunks(data_prefix): @@ -54,7 +51,9 @@ def get_latest_chunks(path, num_chunks, allow_less, sort_key_fn): chunks = get_all_chunks(path) if len(chunks) < num_chunks: if allow_less: - print("sorting {} chunks...".format(len(chunks)), end='', flush=True) + print("sorting {} chunks...".format(len(chunks)), + end='', + flush=True) chunks.sort(key=sort_key_fn, reverse=True) print("[done]") print("{} - {}".format(os.path.basename(chunks[-1]), @@ -80,287 +79,31 @@ def identity_function(name): def game_number_for_name(name): - num_str = os.path.basename(name).upper().strip("ABCDEFGHIJKLMNOPQRSTUVWXYZ_-.") + num_str = os.path.basename(name).upper().strip( + "ABCDEFGHIJKLMNOPQRSTUVWXYZ_-.") return int(num_str) -def extract_policy_bits(raw): - # Next 7432 are easy, policy extraction. - policy = tf.io.decode_raw(tf.strings.substr(raw, 8, 7432), tf.float32) - # Next are 104 bit packed chess boards, they have to be expanded. - bit_planes = tf.expand_dims( - tf.reshape( - tf.io.decode_raw(tf.strings.substr(raw, 7440, 832), tf.uint8), - [-1, 104, 8]), -1) - bit_planes = tf.bitwise.bitwise_and(tf.tile(bit_planes, [1, 1, 1, 8]), - [128, 64, 32, 16, 8, 4, 2, 1]) - bit_planes = tf.minimum(1., tf.cast(bit_planes, tf.float32)) - return policy, bit_planes - - -def extract_byte_planes(raw): - # 5 bytes in input are expanded and tiled - unit_planes = tf.expand_dims( - tf.expand_dims( - tf.io.decode_raw(tf.strings.substr(raw, 8272, 5), tf.uint8), -1), - -1) - unit_planes = tf.tile(unit_planes, [1, 1, 8, 8]) - return unit_planes - - -def extract_rule50_zero_one(raw): - # rule50 count plane. - rule50_plane = tf.expand_dims( - tf.expand_dims( - tf.io.decode_raw(tf.strings.substr(raw, 8277, 1), tf.uint8), -1), - -1) - rule50_plane = tf.cast(tf.tile(rule50_plane, [1, 1, 8, 8]), tf.float32) - rule50_plane = tf.divide(rule50_plane, 99.) - # zero plane and one plane - zero_plane = tf.zeros_like(rule50_plane) - one_plane = tf.ones_like(rule50_plane) - return rule50_plane, zero_plane, one_plane - - -def extract_rule50_100_zero_one(raw): - # rule50 count plane. - rule50_plane = tf.expand_dims( - tf.expand_dims( - tf.io.decode_raw(tf.strings.substr(raw, 8277, 1), tf.uint8), -1), - -1) - rule50_plane = tf.cast(tf.tile(rule50_plane, [1, 1, 8, 8]), tf.float32) - rule50_plane = tf.divide(rule50_plane, 100.) - # zero plane and one plane - zero_plane = tf.zeros_like(rule50_plane) - one_plane = tf.ones_like(rule50_plane) - return rule50_plane, zero_plane, one_plane - - -def extract_invariance(raw): - # invariance plane. - invariance_plane = tf.expand_dims( - tf.expand_dims( - tf.io.decode_raw(tf.strings.substr(raw, 8278, 1), tf.uint8), -1), - -1) - return tf.cast(tf.tile(invariance_plane, [1, 1, 8, 8]), tf.float32) - - -def extract_outputs(raw): - # winner is stored in one signed byte and needs to be converted to one hot. - winner = tf.cast( - tf.io.decode_raw(tf.strings.substr(raw, 8279, 1), tf.int8), tf.float32) - winner = tf.tile(winner, [1, 3]) - z = tf.cast(tf.equal(winner, [1., 0., -1.]), tf.float32) - - # Outcome distribution needs to be calculated from q and d. - best_q = tf.io.decode_raw(tf.strings.substr(raw, 8284, 4), tf.float32) - best_d = tf.io.decode_raw(tf.strings.substr(raw, 8292, 4), tf.float32) - best_q_w = 0.5 * (1.0 - best_d + best_q) - best_q_l = 0.5 * (1.0 - best_d - best_q) - - q = tf.concat([best_q_w, best_d, best_q_l], 1) - - ply_count = tf.io.decode_raw(tf.strings.substr(raw, 8304, 4), tf.float32) - return z, q, ply_count - - -def extract_inputs_outputs_if1(raw): - # first 4 bytes in each batch entry are boring. - # Next 4 change how we construct some of the unit planes. - #input_format = tf.reshape( - # tf.io.decode_raw(tf.strings.substr(raw, 4, 4), tf.int32), - # [-1, 1, 1, 1]) - # tf.debugging.assert_equal(input_format, tf.ones_like(input_format)) - - policy, bit_planes = extract_policy_bits(raw) - - # Next 5 are castling + stm, all of which simply copy the byte value to all squares. - unit_planes = tf.cast(extract_byte_planes(raw), tf.float32) - - rule50_plane, zero_plane, one_plane = extract_rule50_zero_one(raw) - - inputs = tf.reshape( - tf.concat( - [bit_planes, unit_planes, rule50_plane, zero_plane, one_plane], 1), - [-1, 112, 64]) - - z, q, ply_count = extract_outputs(raw) - - return (inputs, policy, z, q, ply_count) - - -def extract_unit_planes_with_bitsplat(raw): - unit_planes = extract_byte_planes(raw) - bitsplat_unit_planes = tf.bitwise.bitwise_and( - unit_planes, [1, 2, 4, 8, 16, 32, 64, 128]) - bitsplat_unit_planes = tf.minimum( - 1., tf.cast(bitsplat_unit_planes, tf.float32)) - unit_planes = tf.cast(unit_planes, tf.float32) - return unit_planes, bitsplat_unit_planes - - -def make_frc_castling(bitsplat_unit_planes, zero_plane): - queenside = tf.concat([ - bitsplat_unit_planes[:, :1, :1], zero_plane[:, :, :6], - bitsplat_unit_planes[:, 2:3, :1] - ], 2) - kingside = tf.concat([ - bitsplat_unit_planes[:, 1:2, :1], zero_plane[:, :, :6], - bitsplat_unit_planes[:, 3:4, :1] - ], 2) - return queenside, kingside - - -def extract_inputs_outputs_if2(raw): - # first 4 bytes in each batch entry are boring. - # Next 4 change how we construct some of the unit planes. - #input_format = tf.reshape( - # tf.io.decode_raw(tf.strings.substr(raw, 4, 4), tf.int32), - # [-1, 1, 1, 1]) - # tf.debugging.assert_equal(input_format, tf.multiply(tf.ones_like(input_format), 2)) - - policy, bit_planes = extract_policy_bits(raw) - - # Next 5 inputs are 4 frc castling and 1 stm. - # In order to do frc we need to make bit unpacked versions. Note little endian for these fields so the bitwise_and array is reversed. - # Although we only need bit unpacked for first 4 of 5 planes, its simpler just to create them all. - unit_planes, bitsplat_unit_planes = extract_unit_planes_with_bitsplat(raw) - - rule50_plane, zero_plane, one_plane = extract_rule50_zero_one(raw) - - # For FRC the old unit planes must be replaced with 0 and 2 merged, 1 and 3 merged, two zero planes and then original 4. - queenside, kingside = make_frc_castling(bitsplat_unit_planes, zero_plane) - unit_planes = tf.concat( - [queenside, kingside, zero_plane, zero_plane, unit_planes[:, 4:]], 1) - - inputs = tf.reshape( - tf.concat( - [bit_planes, unit_planes, rule50_plane, zero_plane, one_plane], 1), - [-1, 112, 64]) - - z, q, ply_count = extract_outputs(raw) - - return (inputs, policy, z, q, ply_count) - - -def make_canonical_unit_planes(bitsplat_unit_planes, zero_plane): - # For canonical the old unit planes must be replaced with 0 and 2 merged, 1 and 3 merged, two zero planes and then en-passant. - queenside, kingside = make_frc_castling(bitsplat_unit_planes, zero_plane) - enpassant = tf.concat( - [zero_plane[:, :, :7], bitsplat_unit_planes[:, 4:, :1]], 2) - unit_planes = tf.concat( - [queenside, kingside, zero_plane, zero_plane, enpassant], 1) - return unit_planes - - -def extract_inputs_outputs_if3(raw): - # first 4 bytes in each batch entry are boring. - # Next 4 change how we construct some of the unit planes. - #input_format = tf.reshape( - # tf.io.decode_raw(tf.strings.substr(raw, 4, 4), tf.int32), - # [-1, 1, 1, 1]) - # tf.debugging.assert_equal(input_format, tf.multiply(tf.ones_like(input_format), 3)) - - policy, bit_planes = extract_policy_bits(raw) - - # Next 5 inputs are 4 castling and 1 enpassant. - # In order to do the frc castling and if3 enpassant plane we need to make bit unpacked versions. Note little endian for these fields so the bitwise_and array is reversed. - unit_planes, bitsplat_unit_planes = extract_unit_planes_with_bitsplat(raw) - - rule50_plane, zero_plane, one_plane = extract_rule50_zero_one(raw) - - unit_planes = make_canonical_unit_planes(bitsplat_unit_planes, zero_plane) - - inputs = tf.reshape( - tf.concat( - [bit_planes, unit_planes, rule50_plane, zero_plane, one_plane], 1), - [-1, 112, 64]) - - z, q, ply_count = extract_outputs(raw) - - return (inputs, policy, z, q, ply_count) - - -def extract_inputs_outputs_if4(raw): - # first 4 bytes in each batch entry are boring. - # Next 4 change how we construct some of the unit planes. - #input_format = tf.reshape( - # tf.io.decode_raw(tf.strings.substr(raw, 4, 4), tf.int32), - # [-1, 1, 1, 1]) - # tf.debugging.assert_equal(input_format, tf.multiply(tf.ones_like(input_format), 3)) - - policy, bit_planes = extract_policy_bits(raw) - - # Next 5 inputs are 4 castling and 1 enpassant. - # In order to do the frc castling and if3 enpassant plane we need to make bit unpacked versions. Note little endian for these fields so the bitwise_and array is reversed. - unit_planes, bitsplat_unit_planes = extract_unit_planes_with_bitsplat(raw) - - rule50_plane, zero_plane, one_plane = extract_rule50_100_zero_one(raw) - - unit_planes = make_canonical_unit_planes(bitsplat_unit_planes, zero_plane) - - inputs = tf.reshape( - tf.concat( - [bit_planes, unit_planes, rule50_plane, zero_plane, one_plane], 1), - [-1, 112, 64]) - - z, q, ply_count = extract_outputs(raw) - - return (inputs, policy, z, q, ply_count) - - -def make_armageddon_stm(invariance_plane): - # invariance_plane contains values of 128 or higher if its black side to move, 127 or lower otherwise. - # Convert this to 0,1 by subtracting off 127 and then clipping. - return tf.clip_by_value(invariance_plane - 127., 0., 1.) - - -def extract_inputs_outputs_if132(raw): - # first 4 bytes in each batch entry are boring. - # Next 4 change how we construct some of the unit planes. - #input_format = tf.reshape( - # tf.io.decode_raw(tf.strings.substr(raw, 4, 4), tf.int32), - # [-1, 1, 1, 1]) - # tf.debugging.assert_equal(input_format, tf.multiply(tf.ones_like(input_format), 3)) - - policy, bit_planes = extract_policy_bits(raw) - - # Next 5 inputs are 4 castling and 1 enpassant. - # In order to do the frc castling and if3 enpassant plane we need to make bit unpacked versions. Note little endian for these fields so the bitwise_and array is reversed. - unit_planes, bitsplat_unit_planes = extract_unit_planes_with_bitsplat(raw) - - rule50_plane, zero_plane, one_plane = extract_rule50_100_zero_one(raw) - - unit_planes = make_canonical_unit_planes(bitsplat_unit_planes, zero_plane) - - armageddon_stm = make_armageddon_stm(extract_invariance(raw)) - - inputs = tf.reshape( - tf.concat( - [bit_planes, unit_planes, rule50_plane, armageddon_stm, one_plane], - 1), [-1, 112, 64]) - - z, q, ply_count = extract_outputs(raw) - - return (inputs, policy, z, q, ply_count) - - -def select_extractor(mode): - if mode == 1: - return extract_inputs_outputs_if1 - if mode == 2: - return extract_inputs_outputs_if2 - if mode == 3: - return extract_inputs_outputs_if3 - if mode == 4 or mode == 5: - return extract_inputs_outputs_if4 - if mode == 132 or mode == 133: - return extract_inputs_outputs_if132 - assert (false) - - -def semi_sample(x): - return tf.slice(tf.random.shuffle(x), [0], [SKIP_MULTIPLE]) +def get_input_mode(cfg): + import proto.net_pb2 as pb + input_mode = cfg['model'].get('input_type', 'classic') + + if input_mode == "classic": + return pb.NetworkFormat.INPUT_CLASSICAL_112_PLANE + elif input_mode == "frc_castling": + return pb.NetworkFormat.INPUT_112_WITH_CASTLING_PLANE + elif input_mode == "canonical": + return pb.NetworkFormat.INPUT_112_WITH_CANONICALIZATION + elif input_mode == "canonical_100": + return pb.NetworkFormat.INPUT_112_WITH_CANONICALIZATION_HECTOPLIES + elif input_mode == "canonical_armageddon": + return pb.NetworkFormat.INPUT_112_WITH_CANONICALIZATION_HECTOPLIES_ARMAGEDDON + elif input_mode == "canonical_v2": + return pb.NetworkFormat.INPUT_112_WITH_CANONICALIZATION_V2 + elif input_mode == "canonical_v2_armageddon": + return pb.NetworkFormat.INPUT_112_WITH_CANONICALIZATION_V2_ARMAGEDDON + else: + raise ValueError("Unknown input mode format: {}".format(input_mode)) def main(cmd): @@ -370,8 +113,6 @@ def main(cmd): num_chunks = cfg['dataset']['num_chunks'] allow_less = cfg['dataset'].get('allow_less_chunks', False) train_ratio = cfg['dataset']['train_ratio'] - experimental_parser = cfg['dataset'].get('experimental_v5_only_dataset', - False) num_train = int(num_chunks * train_ratio) num_test = num_chunks - num_train sort_type = cfg['dataset'].get('sort_type', 'mtime') @@ -405,83 +146,78 @@ def main(cmd): if total_batch_size % batch_splits != 0: raise ValueError('num_batch_splits must divide batch_size evenly') split_batch_size = total_batch_size // batch_splits - # Load data with split batch size, which will be combined to the total batch size in tfprocess. - ChunkParser.BATCH_SIZE = split_batch_size - value_focus_min = cfg['training'].get('value_focus_min', 1) - value_focus_slope = cfg['training'].get('value_focus_slope', 0) + diff_focus_min = cfg['training'].get('diff_focus_min', 1) + diff_focus_slope = cfg['training'].get('diff_focus_slope', 0) + diff_focus_q_weight = cfg['training'].get('diff_focus_q_weight', 6.0) + diff_focus_pol_scale = cfg['training'].get('diff_focus_pol_scale', 3.5) root_dir = os.path.join(cfg['training']['path'], cfg['name']) if not os.path.exists(root_dir): os.makedirs(root_dir) - tfprocess = TFProcess(cfg) - experimental_reads = max(2, mp.cpu_count() - 2) // 2 - extractor = select_extractor(tfprocess.INPUT_MODE) - - if experimental_parser and (value_focus_min != 1 or value_focus_slope != 0): - raise ValueError('Experimental parser does not support non-default value \ - focus parameters.') - - def read(x): - return tf.data.FixedLengthRecordDataset( - x, - 8308, - compression_type='GZIP', - num_parallel_reads=experimental_reads) - - if experimental_parser: - train_dataset = tf.data.Dataset.from_tensor_slices(train_chunks).shuffle(len(train_chunks)).repeat().batch(256)\ - .interleave(read, num_parallel_calls=2)\ - .batch(SKIP_MULTIPLE*SKIP).map(semi_sample).unbatch()\ - .shuffle(shuffle_size)\ - .batch(split_batch_size).map(extractor).prefetch(4) - else: - train_parser = ChunkParser(train_chunks, - tfprocess.INPUT_MODE, - shuffle_size=shuffle_size, - sample=SKIP, - batch_size=ChunkParser.BATCH_SIZE, - value_focus_min=value_focus_min, - value_focus_slope=value_focus_slope, - workers=train_workers) - train_dataset = tf.data.Dataset.from_generator( - train_parser.parse, - output_types=(tf.string, tf.string, tf.string, tf.string, - tf.string)) - train_dataset = train_dataset.map(ChunkParser.parse_function) - train_dataset = train_dataset.prefetch(4) - shuffle_size = int(shuffle_size * (1.0 - train_ratio)) - if experimental_parser: - test_dataset = tf.data.Dataset.from_tensor_slices(test_chunks).shuffle(len(test_chunks)).repeat().batch(256)\ - .interleave(read, num_parallel_calls=2)\ - .batch(SKIP_MULTIPLE*SKIP).map(semi_sample).unbatch()\ - .shuffle(shuffle_size)\ - .batch(split_batch_size).map(extractor).prefetch(4) - else: - # no value focus for test_parser - test_parser = ChunkParser(test_chunks, - tfprocess.INPUT_MODE, - shuffle_size=shuffle_size, - sample=SKIP, - batch_size=ChunkParser.BATCH_SIZE, - workers=test_workers) - test_dataset = tf.data.Dataset.from_generator( - test_parser.parse, - output_types=(tf.string, tf.string, tf.string, tf.string, - tf.string)) - test_dataset = test_dataset.map(ChunkParser.parse_function) - test_dataset = test_dataset.prefetch(4) + train_parser = ChunkParser(train_chunks, + get_input_mode(cfg), + shuffle_size=shuffle_size, + sample=SKIP, + batch_size=split_batch_size, + diff_focus_min=diff_focus_min, + diff_focus_slope=diff_focus_slope, + diff_focus_q_weight=diff_focus_q_weight, + diff_focus_pol_scale=diff_focus_pol_scale, + workers=train_workers) + test_shuffle_size = int(shuffle_size * (1.0 - train_ratio)) + # no diff focus for test_parser + test_parser = ChunkParser(test_chunks, + get_input_mode(cfg), + shuffle_size=test_shuffle_size, + sample=SKIP, + batch_size=split_batch_size, + workers=test_workers) + if 'input_validation' in cfg['dataset']: + valid_chunks = get_all_chunks(cfg['dataset']['input_validation']) + validation_parser = ChunkParser(valid_chunks, + get_input_mode(cfg), + sample=1, + batch_size=split_batch_size, + workers=0) + + import tensorflow as tf + from chunkparsefunc import parse_function + from tfprocess import TFProcess + tfprocess = TFProcess(cfg) + train_dataset = tf.data.Dataset.from_generator( + train_parser.parse, + output_types=(tf.string, tf.string, tf.string, tf.string, tf.string)) + train_dataset = train_dataset.map(parse_function) + test_dataset = tf.data.Dataset.from_generator( + test_parser.parse, + output_types=(tf.string, tf.string, tf.string, tf.string, tf.string)) + test_dataset = test_dataset.map(parse_function) validation_dataset = None if 'input_validation' in cfg['dataset']: - valid_chunks = get_all_chunks(cfg['dataset']['input_validation']) - validation_dataset = tf.data.FixedLengthRecordDataset(valid_chunks, 8308, compression_type='GZIP', num_parallel_reads=experimental_reads)\ - .batch(split_batch_size, drop_remainder=True).map(extractor).prefetch(4) + validation_dataset = tf.data.Dataset.from_generator( + validation_parser.sequential, + output_types=(tf.string, tf.string, tf.string, tf.string, + tf.string)) + validation_dataset = validation_dataset.map(parse_function) - tfprocess.init_v2(train_dataset, test_dataset, validation_dataset) + if tfprocess.strategy is None: #Mirrored strategy appends prefetch itself with a value depending on number of replicas + train_dataset = train_dataset.prefetch(4) + test_dataset = test_dataset.prefetch(4) + if validation_dataset is not None: + validation_dataset = validation_dataset.prefetch(4) + else: + options = tf.data.Options() + options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF + train_dataset = train_dataset.with_options(options) + test_dataset = test_dataset.with_options(options) + if validation_dataset is not None: + validation_dataset = validation_dataset.with_options(options) + tfprocess.init(train_dataset, test_dataset, validation_dataset) - tfprocess.restore_v2() + tfprocess.restore() # If number of test positions is not given # sweeps through all test chunks statistically @@ -490,18 +226,18 @@ def read(x): # This does not affect results, because test results are simple averages that are independent of batch size. num_evals = cfg['training'].get('num_test_positions', len(test_chunks) * 10) - num_evals = max(1, num_evals // ChunkParser.BATCH_SIZE) + num_evals = max(1, num_evals // split_batch_size) print("Using {} evaluation batches".format(num_evals)) - - tfprocess.process_loop_v2(total_batch_size, - num_evals, - batch_splits=batch_splits) + tfprocess.total_batch_size = total_batch_size + tfprocess.process_loop(total_batch_size, + num_evals, + batch_splits=batch_splits) if cmd.output is not None: if cfg['training'].get('swa_output', False): - tfprocess.save_swa_weights_v2(cmd.output) + tfprocess.save_swa_weights(cmd.output) else: - tfprocess.save_leelaz_weights_v2(cmd.output) + tfprocess.save_leelaz_weights(cmd.output) train_parser.shutdown() test_parser.shutdown() diff --git a/tf/update_steps.py b/tf/update_steps.py index 0477a101..953dfcab 100644 --- a/tf/update_steps.py +++ b/tf/update_steps.py @@ -18,9 +18,9 @@ def main(cmd): os.makedirs(root_dir) tfprocess = TFProcess(cfg) - tfprocess.init_net_v2() + tfprocess.init_net() - tfprocess.restore_v2() + tfprocess.restore() START_FROM = cmd.start diff --git a/uv.lock b/uv.lock new file mode 100644 index 00000000..f8d4106f --- /dev/null +++ b/uv.lock @@ -0,0 +1,2287 @@ +version = 1 +revision = 3 +requires-python = ">=3.13" + +[[package]] +name = "absl-py" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/64/c7/8de93764ad66968d19329a7e0c147a2bb3c7054c554d4a119111b8f9440f/absl_py-2.4.0.tar.gz", hash = "sha256:8c6af82722b35cf71e0f4d1d47dcaebfff286e27110a99fc359349b247dfb5d4", size = 116543, upload-time = "2026-01-28T10:17:05.322Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/a6/907a406bb7d359e6a63f99c313846d9eec4f7e6f7437809e03aa00fa3074/absl_py-2.4.0-py3-none-any.whl", hash = "sha256:88476fd881ca8aab94ffa78b7b6c632a782ab3ba1cd19c9bd423abc4fb4cd28d", size = 135750, upload-time = "2026-01-28T10:17:04.19Z" }, +] + +[[package]] +name = "aiofiles" +version = "25.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/41/c3/534eac40372d8ee36ef40df62ec129bee4fdb5ad9706e58a29be53b2c970/aiofiles-25.1.0.tar.gz", hash = "sha256:a8d728f0a29de45dc521f18f07297428d56992a742f0cd2701ba86e44d23d5b2", size = 46354, upload-time = "2025-10-09T20:51:04.358Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/8a/340a1555ae33d7354dbca4faa54948d76d89a27ceef032c8c3bc661d003e/aiofiles-25.1.0-py3-none-any.whl", hash = "sha256:abe311e527c862958650f9438e859c1fa7568a141b22abcd015e120e86a85695", size = 14668, upload-time = "2025-10-09T20:51:03.174Z" }, +] + +[[package]] +name = "aiohappyeyeballs" +version = "2.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/26/30/f84a107a9c4331c14b2b586036f40965c128aa4fee4dda5d3d51cb14ad54/aiohappyeyeballs-2.6.1.tar.gz", hash = "sha256:c3f9d0113123803ccadfdf3f0faa505bc78e6a72d1cc4806cbd719826e943558", size = 22760, upload-time = "2025-03-12T01:42:48.764Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/15/5bf3b99495fb160b63f95972b81750f18f7f4e02ad051373b669d17d44f2/aiohappyeyeballs-2.6.1-py3-none-any.whl", hash = "sha256:f349ba8f4b75cb25c99c5c2d84e997e485204d2902a9597802b0371f09331fb8", size = 15265, upload-time = "2025-03-12T01:42:47.083Z" }, +] + +[[package]] +name = "aiohttp" +version = "3.13.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohappyeyeballs" }, + { name = "aiosignal" }, + { name = "attrs" }, + { name = "frozenlist" }, + { name = "multidict" }, + { name = "propcache" }, + { name = "yarl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/50/42/32cf8e7704ceb4481406eb87161349abb46a57fee3f008ba9cb610968646/aiohttp-3.13.3.tar.gz", hash = "sha256:a949eee43d3782f2daae4f4a2819b2cb9b0c5d3b7f7a927067cc84dafdbb9f88", size = 7844556, upload-time = "2026-01-03T17:33:05.204Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/97/8a/12ca489246ca1faaf5432844adbfce7ff2cc4997733e0af120869345643a/aiohttp-3.13.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:5dff64413671b0d3e7d5918ea490bdccb97a4ad29b3f311ed423200b2203e01c", size = 734190, upload-time = "2026-01-03T17:30:45.832Z" }, + { url = "https://files.pythonhosted.org/packages/32/08/de43984c74ed1fca5c014808963cc83cb00d7bb06af228f132d33862ca76/aiohttp-3.13.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:87b9aab6d6ed88235aa2970294f496ff1a1f9adcd724d800e9b952395a80ffd9", size = 491783, upload-time = "2026-01-03T17:30:47.466Z" }, + { url = "https://files.pythonhosted.org/packages/17/f8/8dd2cf6112a5a76f81f81a5130c57ca829d101ad583ce57f889179accdda/aiohttp-3.13.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:425c126c0dc43861e22cb1c14ba4c8e45d09516d0a3ae0a3f7494b79f5f233a3", size = 490704, upload-time = "2026-01-03T17:30:49.373Z" }, + { url = "https://files.pythonhosted.org/packages/6d/40/a46b03ca03936f832bc7eaa47cfbb1ad012ba1be4790122ee4f4f8cba074/aiohttp-3.13.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7f9120f7093c2a32d9647abcaf21e6ad275b4fbec5b55969f978b1a97c7c86bf", size = 1720652, upload-time = "2026-01-03T17:30:50.974Z" }, + { url = "https://files.pythonhosted.org/packages/f7/7e/917fe18e3607af92657e4285498f500dca797ff8c918bd7d90b05abf6c2a/aiohttp-3.13.3-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:697753042d57f4bf7122cab985bf15d0cef23c770864580f5af4f52023a56bd6", size = 1692014, upload-time = "2026-01-03T17:30:52.729Z" }, + { url = "https://files.pythonhosted.org/packages/71/b6/cefa4cbc00d315d68973b671cf105b21a609c12b82d52e5d0c9ae61d2a09/aiohttp-3.13.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6de499a1a44e7de70735d0b39f67c8f25eb3d91eb3103be99ca0fa882cdd987d", size = 1759777, upload-time = "2026-01-03T17:30:54.537Z" }, + { url = "https://files.pythonhosted.org/packages/fb/e3/e06ee07b45e59e6d81498b591fc589629be1553abb2a82ce33efe2a7b068/aiohttp-3.13.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:37239e9f9a7ea9ac5bf6b92b0260b01f8a22281996da609206a84df860bc1261", size = 1861276, upload-time = "2026-01-03T17:30:56.512Z" }, + { url = "https://files.pythonhosted.org/packages/7c/24/75d274228acf35ceeb2850b8ce04de9dd7355ff7a0b49d607ee60c29c518/aiohttp-3.13.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f76c1e3fe7d7c8afad7ed193f89a292e1999608170dcc9751a7462a87dfd5bc0", size = 1743131, upload-time = "2026-01-03T17:30:58.256Z" }, + { url = "https://files.pythonhosted.org/packages/04/98/3d21dde21889b17ca2eea54fdcff21b27b93f45b7bb94ca029c31ab59dc3/aiohttp-3.13.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:fc290605db2a917f6e81b0e1e0796469871f5af381ce15c604a3c5c7e51cb730", size = 1556863, upload-time = "2026-01-03T17:31:00.445Z" }, + { url = "https://files.pythonhosted.org/packages/9e/84/da0c3ab1192eaf64782b03971ab4055b475d0db07b17eff925e8c93b3aa5/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4021b51936308aeea0367b8f006dc999ca02bc118a0cc78c303f50a2ff6afb91", size = 1682793, upload-time = "2026-01-03T17:31:03.024Z" }, + { url = "https://files.pythonhosted.org/packages/ff/0f/5802ada182f575afa02cbd0ec5180d7e13a402afb7c2c03a9aa5e5d49060/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:49a03727c1bba9a97d3e93c9f93ca03a57300f484b6e935463099841261195d3", size = 1716676, upload-time = "2026-01-03T17:31:04.842Z" }, + { url = "https://files.pythonhosted.org/packages/3f/8c/714d53bd8b5a4560667f7bbbb06b20c2382f9c7847d198370ec6526af39c/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:3d9908a48eb7416dc1f4524e69f1d32e5d90e3981e4e37eb0aa1cd18f9cfa2a4", size = 1733217, upload-time = "2026-01-03T17:31:06.868Z" }, + { url = "https://files.pythonhosted.org/packages/7d/79/e2176f46d2e963facea939f5be2d26368ce543622be6f00a12844d3c991f/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:2712039939ec963c237286113c68dbad80a82a4281543f3abf766d9d73228998", size = 1552303, upload-time = "2026-01-03T17:31:08.958Z" }, + { url = "https://files.pythonhosted.org/packages/ab/6a/28ed4dea1759916090587d1fe57087b03e6c784a642b85ef48217b0277ae/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:7bfdc049127717581866fa4708791220970ce291c23e28ccf3922c700740fdc0", size = 1763673, upload-time = "2026-01-03T17:31:10.676Z" }, + { url = "https://files.pythonhosted.org/packages/e8/35/4a3daeb8b9fab49240d21c04d50732313295e4bd813a465d840236dd0ce1/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8057c98e0c8472d8846b9c79f56766bcc57e3e8ac7bfd510482332366c56c591", size = 1721120, upload-time = "2026-01-03T17:31:12.575Z" }, + { url = "https://files.pythonhosted.org/packages/bc/9f/d643bb3c5fb99547323e635e251c609fbbc660d983144cfebec529e09264/aiohttp-3.13.3-cp313-cp313-win32.whl", hash = "sha256:1449ceddcdbcf2e0446957863af03ebaaa03f94c090f945411b61269e2cb5daf", size = 427383, upload-time = "2026-01-03T17:31:14.382Z" }, + { url = "https://files.pythonhosted.org/packages/4e/f1/ab0395f8a79933577cdd996dd2f9aa6014af9535f65dddcf88204682fe62/aiohttp-3.13.3-cp313-cp313-win_amd64.whl", hash = "sha256:693781c45a4033d31d4187d2436f5ac701e7bbfe5df40d917736108c1cc7436e", size = 453899, upload-time = "2026-01-03T17:31:15.958Z" }, + { url = "https://files.pythonhosted.org/packages/99/36/5b6514a9f5d66f4e2597e40dea2e3db271e023eb7a5d22defe96ba560996/aiohttp-3.13.3-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:ea37047c6b367fd4bd632bff8077449b8fa034b69e812a18e0132a00fae6e808", size = 737238, upload-time = "2026-01-03T17:31:17.909Z" }, + { url = "https://files.pythonhosted.org/packages/f7/49/459327f0d5bcd8c6c9ca69e60fdeebc3622861e696490d8674a6d0cb90a6/aiohttp-3.13.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:6fc0e2337d1a4c3e6acafda6a78a39d4c14caea625124817420abceed36e2415", size = 492292, upload-time = "2026-01-03T17:31:19.919Z" }, + { url = "https://files.pythonhosted.org/packages/e8/0b/b97660c5fd05d3495b4eb27f2d0ef18dc1dc4eff7511a9bf371397ff0264/aiohttp-3.13.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c685f2d80bb67ca8c3837823ad76196b3694b0159d232206d1e461d3d434666f", size = 493021, upload-time = "2026-01-03T17:31:21.636Z" }, + { url = "https://files.pythonhosted.org/packages/54/d4/438efabdf74e30aeceb890c3290bbaa449780583b1270b00661126b8aae4/aiohttp-3.13.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:48e377758516d262bde50c2584fc6c578af272559c409eecbdd2bae1601184d6", size = 1717263, upload-time = "2026-01-03T17:31:23.296Z" }, + { url = "https://files.pythonhosted.org/packages/71/f2/7bddc7fd612367d1459c5bcf598a9e8f7092d6580d98de0e057eb42697ad/aiohttp-3.13.3-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:34749271508078b261c4abb1767d42b8d0c0cc9449c73a4df494777dc55f0687", size = 1669107, upload-time = "2026-01-03T17:31:25.334Z" }, + { url = "https://files.pythonhosted.org/packages/00/5a/1aeaecca40e22560f97610a329e0e5efef5e0b5afdf9f857f0d93839ab2e/aiohttp-3.13.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:82611aeec80eb144416956ec85b6ca45a64d76429c1ed46ae1b5f86c6e0c9a26", size = 1760196, upload-time = "2026-01-03T17:31:27.394Z" }, + { url = "https://files.pythonhosted.org/packages/f8/f8/0ff6992bea7bd560fc510ea1c815f87eedd745fe035589c71ce05612a19a/aiohttp-3.13.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2fff83cfc93f18f215896e3a190e8e5cb413ce01553901aca925176e7568963a", size = 1843591, upload-time = "2026-01-03T17:31:29.238Z" }, + { url = "https://files.pythonhosted.org/packages/e3/d1/e30e537a15f53485b61f5be525f2157da719819e8377298502aebac45536/aiohttp-3.13.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bbe7d4cecacb439e2e2a8a1a7b935c25b812af7a5fd26503a66dadf428e79ec1", size = 1720277, upload-time = "2026-01-03T17:31:31.053Z" }, + { url = "https://files.pythonhosted.org/packages/84/45/23f4c451d8192f553d38d838831ebbc156907ea6e05557f39563101b7717/aiohttp-3.13.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b928f30fe49574253644b1ca44b1b8adbd903aa0da4b9054a6c20fc7f4092a25", size = 1548575, upload-time = "2026-01-03T17:31:32.87Z" }, + { url = "https://files.pythonhosted.org/packages/6a/ed/0a42b127a43712eda7807e7892c083eadfaf8429ca8fb619662a530a3aab/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7b5e8fe4de30df199155baaf64f2fcd604f4c678ed20910db8e2c66dc4b11603", size = 1679455, upload-time = "2026-01-03T17:31:34.76Z" }, + { url = "https://files.pythonhosted.org/packages/2e/b5/c05f0c2b4b4fe2c9d55e73b6d3ed4fd6c9dc2684b1d81cbdf77e7fad9adb/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:8542f41a62bcc58fc7f11cf7c90e0ec324ce44950003feb70640fc2a9092c32a", size = 1687417, upload-time = "2026-01-03T17:31:36.699Z" }, + { url = "https://files.pythonhosted.org/packages/c9/6b/915bc5dad66aef602b9e459b5a973529304d4e89ca86999d9d75d80cbd0b/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:5e1d8c8b8f1d91cd08d8f4a3c2b067bfca6ec043d3ff36de0f3a715feeedf926", size = 1729968, upload-time = "2026-01-03T17:31:38.622Z" }, + { url = "https://files.pythonhosted.org/packages/11/3b/e84581290a9520024a08640b63d07673057aec5ca548177a82026187ba73/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:90455115e5da1c3c51ab619ac57f877da8fd6d73c05aacd125c5ae9819582aba", size = 1545690, upload-time = "2026-01-03T17:31:40.57Z" }, + { url = "https://files.pythonhosted.org/packages/f5/04/0c3655a566c43fd647c81b895dfe361b9f9ad6d58c19309d45cff52d6c3b/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:042e9e0bcb5fba81886c8b4fbb9a09d6b8a00245fd8d88e4d989c1f96c74164c", size = 1746390, upload-time = "2026-01-03T17:31:42.857Z" }, + { url = "https://files.pythonhosted.org/packages/1f/53/71165b26978f719c3419381514c9690bd5980e764a09440a10bb816ea4ab/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2eb752b102b12a76ca02dff751a801f028b4ffbbc478840b473597fc91a9ed43", size = 1702188, upload-time = "2026-01-03T17:31:44.984Z" }, + { url = "https://files.pythonhosted.org/packages/29/a7/cbe6c9e8e136314fa1980da388a59d2f35f35395948a08b6747baebb6aa6/aiohttp-3.13.3-cp314-cp314-win32.whl", hash = "sha256:b556c85915d8efaed322bf1bdae9486aa0f3f764195a0fb6ee962e5c71ef5ce1", size = 433126, upload-time = "2026-01-03T17:31:47.463Z" }, + { url = "https://files.pythonhosted.org/packages/de/56/982704adea7d3b16614fc5936014e9af85c0e34b58f9046655817f04306e/aiohttp-3.13.3-cp314-cp314-win_amd64.whl", hash = "sha256:9bf9f7a65e7aa20dd764151fb3d616c81088f91f8df39c3893a536e279b4b984", size = 459128, upload-time = "2026-01-03T17:31:49.2Z" }, + { url = "https://files.pythonhosted.org/packages/6c/2a/3c79b638a9c3d4658d345339d22070241ea341ed4e07b5ac60fb0f418003/aiohttp-3.13.3-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:05861afbbec40650d8a07ea324367cb93e9e8cc7762e04dd4405df99fa65159c", size = 769512, upload-time = "2026-01-03T17:31:51.134Z" }, + { url = "https://files.pythonhosted.org/packages/29/b9/3e5014d46c0ab0db8707e0ac2711ed28c4da0218c358a4e7c17bae0d8722/aiohttp-3.13.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:2fc82186fadc4a8316768d61f3722c230e2c1dcab4200d52d2ebdf2482e47592", size = 506444, upload-time = "2026-01-03T17:31:52.85Z" }, + { url = "https://files.pythonhosted.org/packages/90/03/c1d4ef9a054e151cd7839cdc497f2638f00b93cbe8043983986630d7a80c/aiohttp-3.13.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:0add0900ff220d1d5c5ebbf99ed88b0c1bbf87aa7e4262300ed1376a6b13414f", size = 510798, upload-time = "2026-01-03T17:31:54.91Z" }, + { url = "https://files.pythonhosted.org/packages/ea/76/8c1e5abbfe8e127c893fe7ead569148a4d5a799f7cf958d8c09f3eedf097/aiohttp-3.13.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:568f416a4072fbfae453dcf9a99194bbb8bdeab718e08ee13dfa2ba0e4bebf29", size = 1868835, upload-time = "2026-01-03T17:31:56.733Z" }, + { url = "https://files.pythonhosted.org/packages/8e/ac/984c5a6f74c363b01ff97adc96a3976d9c98940b8969a1881575b279ac5d/aiohttp-3.13.3-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:add1da70de90a2569c5e15249ff76a631ccacfe198375eead4aadf3b8dc849dc", size = 1720486, upload-time = "2026-01-03T17:31:58.65Z" }, + { url = "https://files.pythonhosted.org/packages/b2/9a/b7039c5f099c4eb632138728828b33428585031a1e658d693d41d07d89d1/aiohttp-3.13.3-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:10b47b7ba335d2e9b1239fa571131a87e2d8ec96b333e68b2a305e7a98b0bae2", size = 1847951, upload-time = "2026-01-03T17:32:00.989Z" }, + { url = "https://files.pythonhosted.org/packages/3c/02/3bec2b9a1ba3c19ff89a43a19324202b8eb187ca1e928d8bdac9bbdddebd/aiohttp-3.13.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3dd4dce1c718e38081c8f35f323209d4c1df7d4db4bab1b5c88a6b4d12b74587", size = 1941001, upload-time = "2026-01-03T17:32:03.122Z" }, + { url = "https://files.pythonhosted.org/packages/37/df/d879401cedeef27ac4717f6426c8c36c3091c6e9f08a9178cc87549c537f/aiohttp-3.13.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:34bac00a67a812570d4a460447e1e9e06fae622946955f939051e7cc895cfab8", size = 1797246, upload-time = "2026-01-03T17:32:05.255Z" }, + { url = "https://files.pythonhosted.org/packages/8d/15/be122de1f67e6953add23335c8ece6d314ab67c8bebb3f181063010795a7/aiohttp-3.13.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a19884d2ee70b06d9204b2727a7b9f983d0c684c650254679e716b0b77920632", size = 1627131, upload-time = "2026-01-03T17:32:07.607Z" }, + { url = "https://files.pythonhosted.org/packages/12/12/70eedcac9134cfa3219ab7af31ea56bc877395b1ac30d65b1bc4b27d0438/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5f8ca7f2bb6ba8348a3614c7918cc4bb73268c5ac2a207576b7afea19d3d9f64", size = 1795196, upload-time = "2026-01-03T17:32:09.59Z" }, + { url = "https://files.pythonhosted.org/packages/32/11/b30e1b1cd1f3054af86ebe60df96989c6a414dd87e27ad16950eee420bea/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:b0d95340658b9d2f11d9697f59b3814a9d3bb4b7a7c20b131df4bcef464037c0", size = 1782841, upload-time = "2026-01-03T17:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/88/0d/d98a9367b38912384a17e287850f5695c528cff0f14f791ce8ee2e4f7796/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:a1e53262fd202e4b40b70c3aff944a8155059beedc8a89bba9dc1f9ef06a1b56", size = 1795193, upload-time = "2026-01-03T17:32:13.705Z" }, + { url = "https://files.pythonhosted.org/packages/43/a5/a2dfd1f5ff5581632c7f6a30e1744deda03808974f94f6534241ef60c751/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:d60ac9663f44168038586cab2157e122e46bdef09e9368b37f2d82d354c23f72", size = 1621979, upload-time = "2026-01-03T17:32:15.965Z" }, + { url = "https://files.pythonhosted.org/packages/fa/f0/12973c382ae7c1cccbc4417e129c5bf54c374dfb85af70893646e1f0e749/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:90751b8eed69435bac9ff4e3d2f6b3af1f57e37ecb0fbeee59c0174c9e2d41df", size = 1822193, upload-time = "2026-01-03T17:32:18.219Z" }, + { url = "https://files.pythonhosted.org/packages/3c/5f/24155e30ba7f8c96918af1350eb0663e2430aad9e001c0489d89cd708ab1/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:fc353029f176fd2b3ec6cfc71be166aba1936fe5d73dd1992ce289ca6647a9aa", size = 1769801, upload-time = "2026-01-03T17:32:20.25Z" }, + { url = "https://files.pythonhosted.org/packages/eb/f8/7314031ff5c10e6ece114da79b338ec17eeff3a079e53151f7e9f43c4723/aiohttp-3.13.3-cp314-cp314t-win32.whl", hash = "sha256:2e41b18a58da1e474a057b3d35248d8320029f61d70a37629535b16a0c8f3767", size = 466523, upload-time = "2026-01-03T17:32:22.215Z" }, + { url = "https://files.pythonhosted.org/packages/b4/63/278a98c715ae467624eafe375542d8ba9b4383a016df8fdefe0ae28382a7/aiohttp-3.13.3-cp314-cp314t-win_amd64.whl", hash = "sha256:44531a36aa2264a1860089ffd4dce7baf875ee5a6079d5fb42e261c704ef7344", size = 499694, upload-time = "2026-01-03T17:32:24.546Z" }, +] + +[[package]] +name = "aiohttp-jinja2" +version = "1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "jinja2" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e6/39/da5a94dd89b1af7241fb7fc99ae4e73505b5f898b540b6aba6dc7afe600e/aiohttp-jinja2-1.6.tar.gz", hash = "sha256:a3a7ff5264e5bca52e8ae547bbfd0761b72495230d438d05b6c0915be619b0e2", size = 53057, upload-time = "2023-11-18T15:30:52.559Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/eb/90/65238d4246307195411b87a07d03539049819b022c01bcc773826f600138/aiohttp_jinja2-1.6-py3-none-any.whl", hash = "sha256:0df405ee6ad1b58e5a068a105407dc7dcc1704544c559f1938babde954f945c7", size = 11736, upload-time = "2023-11-18T15:30:50.743Z" }, +] + +[[package]] +name = "aiosignal" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "frozenlist" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, +] + +[[package]] +name = "anyio" +version = "4.12.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "idna" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/96/f0/5eb65b2bb0d09ac6776f2eb54adee6abe8228ea05b20a5ad0e4945de8aac/anyio-4.12.1.tar.gz", hash = "sha256:41cfcc3a4c85d3f05c932da7c26d0201ac36f72abd4435ba90d0464a3ffed703", size = 228685, upload-time = "2026-01-06T11:45:21.246Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/0e/27be9fdef66e72d64c0cdc3cc2823101b80585f8119b5c112c2e8f5f7dab/anyio-4.12.1-py3-none-any.whl", hash = "sha256:d405828884fc140aa80a3c667b8beed277f1dfedec42ba031bd6ac3db606ab6c", size = 113592, upload-time = "2026-01-06T11:45:19.497Z" }, +] + +[[package]] +name = "attrs" +version = "25.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6b/5c/685e6633917e101e5dcb62b9dd76946cbb57c26e133bae9e0cd36033c0a9/attrs-25.4.0.tar.gz", hash = "sha256:16d5969b87f0859ef33a48b35d55ac1be6e42ae49d5e853b597db70c35c57e11", size = 934251, upload-time = "2025-10-06T13:54:44.725Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/2a/7cc015f5b9f5db42b7d48157e23356022889fc354a2813c15934b7cb5c0e/attrs-25.4.0-py3-none-any.whl", hash = "sha256:adcf7e2a1fb3b36ac48d97835bb6d8ade15b8dcce26aba8bf1d14847b57a3373", size = 67615, upload-time = "2025-10-06T13:54:43.17Z" }, +] + +[[package]] +name = "certifi" +version = "2026.1.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/2d/a891ca51311197f6ad14a7ef42e2399f36cf2f9bd44752b3dc4eab60fdc5/certifi-2026.1.4.tar.gz", hash = "sha256:ac726dd470482006e014ad384921ed6438c457018f4b3d204aea4281258b2120", size = 154268, upload-time = "2026-01-04T02:42:41.825Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/ad/3cc14f097111b4de0040c83a525973216457bbeeb63739ef1ed275c1c021/certifi-2026.1.4-py3-none-any.whl", hash = "sha256:9943707519e4add1115f44c2bc244f782c0249876bf51b6599fee1ffbedd685c", size = 152900, upload-time = "2026-01-04T02:42:40.15Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/13/69/33ddede1939fdd074bce5434295f38fae7136463422fe4fd3e0e89b98062/charset_normalizer-3.4.4.tar.gz", hash = "sha256:94537985111c35f28720e43603b8e7b43a6ecfb2ce1d3058bbe955b73404e21a", size = 129418, upload-time = "2025-10-14T04:42:32.879Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/97/45/4b3a1239bbacd321068ea6e7ac28875b03ab8bc0aa0966452db17cd36714/charset_normalizer-3.4.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:e1f185f86a6f3403aa2420e815904c67b2f9ebc443f045edd0de921108345794", size = 208091, upload-time = "2025-10-14T04:41:13.346Z" }, + { url = "https://files.pythonhosted.org/packages/7d/62/73a6d7450829655a35bb88a88fca7d736f9882a27eacdca2c6d505b57e2e/charset_normalizer-3.4.4-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6b39f987ae8ccdf0d2642338faf2abb1862340facc796048b604ef14919e55ed", size = 147936, upload-time = "2025-10-14T04:41:14.461Z" }, + { url = "https://files.pythonhosted.org/packages/89/c5/adb8c8b3d6625bef6d88b251bbb0d95f8205831b987631ab0c8bb5d937c2/charset_normalizer-3.4.4-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3162d5d8ce1bb98dd51af660f2121c55d0fa541b46dff7bb9b9f86ea1d87de72", size = 144180, upload-time = "2025-10-14T04:41:15.588Z" }, + { url = "https://files.pythonhosted.org/packages/91/ed/9706e4070682d1cc219050b6048bfd293ccf67b3d4f5a4f39207453d4b99/charset_normalizer-3.4.4-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:81d5eb2a312700f4ecaa977a8235b634ce853200e828fbadf3a9c50bab278328", size = 161346, upload-time = "2025-10-14T04:41:16.738Z" }, + { url = "https://files.pythonhosted.org/packages/d5/0d/031f0d95e4972901a2f6f09ef055751805ff541511dc1252ba3ca1f80cf5/charset_normalizer-3.4.4-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5bd2293095d766545ec1a8f612559f6b40abc0eb18bb2f5d1171872d34036ede", size = 158874, upload-time = "2025-10-14T04:41:17.923Z" }, + { url = "https://files.pythonhosted.org/packages/f5/83/6ab5883f57c9c801ce5e5677242328aa45592be8a00644310a008d04f922/charset_normalizer-3.4.4-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a8a8b89589086a25749f471e6a900d3f662d1d3b6e2e59dcecf787b1cc3a1894", size = 153076, upload-time = "2025-10-14T04:41:19.106Z" }, + { url = "https://files.pythonhosted.org/packages/75/1e/5ff781ddf5260e387d6419959ee89ef13878229732732ee73cdae01800f2/charset_normalizer-3.4.4-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bc7637e2f80d8530ee4a78e878bce464f70087ce73cf7c1caf142416923b98f1", size = 150601, upload-time = "2025-10-14T04:41:20.245Z" }, + { url = "https://files.pythonhosted.org/packages/d7/57/71be810965493d3510a6ca79b90c19e48696fb1ff964da319334b12677f0/charset_normalizer-3.4.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f8bf04158c6b607d747e93949aa60618b61312fe647a6369f88ce2ff16043490", size = 150376, upload-time = "2025-10-14T04:41:21.398Z" }, + { url = "https://files.pythonhosted.org/packages/e5/d5/c3d057a78c181d007014feb7e9f2e65905a6c4ef182c0ddf0de2924edd65/charset_normalizer-3.4.4-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:554af85e960429cf30784dd47447d5125aaa3b99a6f0683589dbd27e2f45da44", size = 144825, upload-time = "2025-10-14T04:41:22.583Z" }, + { url = "https://files.pythonhosted.org/packages/e6/8c/d0406294828d4976f275ffbe66f00266c4b3136b7506941d87c00cab5272/charset_normalizer-3.4.4-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:74018750915ee7ad843a774364e13a3db91682f26142baddf775342c3f5b1133", size = 162583, upload-time = "2025-10-14T04:41:23.754Z" }, + { url = "https://files.pythonhosted.org/packages/d7/24/e2aa1f18c8f15c4c0e932d9287b8609dd30ad56dbe41d926bd846e22fb8d/charset_normalizer-3.4.4-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:c0463276121fdee9c49b98908b3a89c39be45d86d1dbaa22957e38f6321d4ce3", size = 150366, upload-time = "2025-10-14T04:41:25.27Z" }, + { url = "https://files.pythonhosted.org/packages/e4/5b/1e6160c7739aad1e2df054300cc618b06bf784a7a164b0f238360721ab86/charset_normalizer-3.4.4-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:362d61fd13843997c1c446760ef36f240cf81d3ebf74ac62652aebaf7838561e", size = 160300, upload-time = "2025-10-14T04:41:26.725Z" }, + { url = "https://files.pythonhosted.org/packages/7a/10/f882167cd207fbdd743e55534d5d9620e095089d176d55cb22d5322f2afd/charset_normalizer-3.4.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9a26f18905b8dd5d685d6d07b0cdf98a79f3c7a918906af7cc143ea2e164c8bc", size = 154465, upload-time = "2025-10-14T04:41:28.322Z" }, + { url = "https://files.pythonhosted.org/packages/89/66/c7a9e1b7429be72123441bfdbaf2bc13faab3f90b933f664db506dea5915/charset_normalizer-3.4.4-cp313-cp313-win32.whl", hash = "sha256:9b35f4c90079ff2e2edc5b26c0c77925e5d2d255c42c74fdb70fb49b172726ac", size = 99404, upload-time = "2025-10-14T04:41:29.95Z" }, + { url = "https://files.pythonhosted.org/packages/c4/26/b9924fa27db384bdcd97ab83b4f0a8058d96ad9626ead570674d5e737d90/charset_normalizer-3.4.4-cp313-cp313-win_amd64.whl", hash = "sha256:b435cba5f4f750aa6c0a0d92c541fb79f69a387c91e61f1795227e4ed9cece14", size = 107092, upload-time = "2025-10-14T04:41:31.188Z" }, + { url = "https://files.pythonhosted.org/packages/af/8f/3ed4bfa0c0c72a7ca17f0380cd9e4dd842b09f664e780c13cff1dcf2ef1b/charset_normalizer-3.4.4-cp313-cp313-win_arm64.whl", hash = "sha256:542d2cee80be6f80247095cc36c418f7bddd14f4a6de45af91dfad36d817bba2", size = 100408, upload-time = "2025-10-14T04:41:32.624Z" }, + { url = "https://files.pythonhosted.org/packages/2a/35/7051599bd493e62411d6ede36fd5af83a38f37c4767b92884df7301db25d/charset_normalizer-3.4.4-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:da3326d9e65ef63a817ecbcc0df6e94463713b754fe293eaa03da99befb9a5bd", size = 207746, upload-time = "2025-10-14T04:41:33.773Z" }, + { url = "https://files.pythonhosted.org/packages/10/9a/97c8d48ef10d6cd4fcead2415523221624bf58bcf68a802721a6bc807c8f/charset_normalizer-3.4.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8af65f14dc14a79b924524b1e7fffe304517b2bff5a58bf64f30b98bbc5079eb", size = 147889, upload-time = "2025-10-14T04:41:34.897Z" }, + { url = "https://files.pythonhosted.org/packages/10/bf/979224a919a1b606c82bd2c5fa49b5c6d5727aa47b4312bb27b1734f53cd/charset_normalizer-3.4.4-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:74664978bb272435107de04e36db5a9735e78232b85b77d45cfb38f758efd33e", size = 143641, upload-time = "2025-10-14T04:41:36.116Z" }, + { url = "https://files.pythonhosted.org/packages/ba/33/0ad65587441fc730dc7bd90e9716b30b4702dc7b617e6ba4997dc8651495/charset_normalizer-3.4.4-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:752944c7ffbfdd10c074dc58ec2d5a8a4cd9493b314d367c14d24c17684ddd14", size = 160779, upload-time = "2025-10-14T04:41:37.229Z" }, + { url = "https://files.pythonhosted.org/packages/67/ed/331d6b249259ee71ddea93f6f2f0a56cfebd46938bde6fcc6f7b9a3d0e09/charset_normalizer-3.4.4-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d1f13550535ad8cff21b8d757a3257963e951d96e20ec82ab44bc64aeb62a191", size = 159035, upload-time = "2025-10-14T04:41:38.368Z" }, + { url = "https://files.pythonhosted.org/packages/67/ff/f6b948ca32e4f2a4576aa129d8bed61f2e0543bf9f5f2b7fc3758ed005c9/charset_normalizer-3.4.4-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ecaae4149d99b1c9e7b88bb03e3221956f68fd6d50be2ef061b2381b61d20838", size = 152542, upload-time = "2025-10-14T04:41:39.862Z" }, + { url = "https://files.pythonhosted.org/packages/16/85/276033dcbcc369eb176594de22728541a925b2632f9716428c851b149e83/charset_normalizer-3.4.4-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:cb6254dc36b47a990e59e1068afacdcd02958bdcce30bb50cc1700a8b9d624a6", size = 149524, upload-time = "2025-10-14T04:41:41.319Z" }, + { url = "https://files.pythonhosted.org/packages/9e/f2/6a2a1f722b6aba37050e626530a46a68f74e63683947a8acff92569f979a/charset_normalizer-3.4.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c8ae8a0f02f57a6e61203a31428fa1d677cbe50c93622b4149d5c0f319c1d19e", size = 150395, upload-time = "2025-10-14T04:41:42.539Z" }, + { url = "https://files.pythonhosted.org/packages/60/bb/2186cb2f2bbaea6338cad15ce23a67f9b0672929744381e28b0592676824/charset_normalizer-3.4.4-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:47cc91b2f4dd2833fddaedd2893006b0106129d4b94fdb6af1f4ce5a9965577c", size = 143680, upload-time = "2025-10-14T04:41:43.661Z" }, + { url = "https://files.pythonhosted.org/packages/7d/a5/bf6f13b772fbb2a90360eb620d52ed8f796f3c5caee8398c3b2eb7b1c60d/charset_normalizer-3.4.4-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:82004af6c302b5d3ab2cfc4cc5f29db16123b1a8417f2e25f9066f91d4411090", size = 162045, upload-time = "2025-10-14T04:41:44.821Z" }, + { url = "https://files.pythonhosted.org/packages/df/c5/d1be898bf0dc3ef9030c3825e5d3b83f2c528d207d246cbabe245966808d/charset_normalizer-3.4.4-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:2b7d8f6c26245217bd2ad053761201e9f9680f8ce52f0fcd8d0755aeae5b2152", size = 149687, upload-time = "2025-10-14T04:41:46.442Z" }, + { url = "https://files.pythonhosted.org/packages/a5/42/90c1f7b9341eef50c8a1cb3f098ac43b0508413f33affd762855f67a410e/charset_normalizer-3.4.4-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:799a7a5e4fb2d5898c60b640fd4981d6a25f1c11790935a44ce38c54e985f828", size = 160014, upload-time = "2025-10-14T04:41:47.631Z" }, + { url = "https://files.pythonhosted.org/packages/76/be/4d3ee471e8145d12795ab655ece37baed0929462a86e72372fd25859047c/charset_normalizer-3.4.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:99ae2cffebb06e6c22bdc25801d7b30f503cc87dbd283479e7b606f70aff57ec", size = 154044, upload-time = "2025-10-14T04:41:48.81Z" }, + { url = "https://files.pythonhosted.org/packages/b0/6f/8f7af07237c34a1defe7defc565a9bc1807762f672c0fde711a4b22bf9c0/charset_normalizer-3.4.4-cp314-cp314-win32.whl", hash = "sha256:f9d332f8c2a2fcbffe1378594431458ddbef721c1769d78e2cbc06280d8155f9", size = 99940, upload-time = "2025-10-14T04:41:49.946Z" }, + { url = "https://files.pythonhosted.org/packages/4b/51/8ade005e5ca5b0d80fb4aff72a3775b325bdc3d27408c8113811a7cbe640/charset_normalizer-3.4.4-cp314-cp314-win_amd64.whl", hash = "sha256:8a6562c3700cce886c5be75ade4a5db4214fda19fede41d9792d100288d8f94c", size = 107104, upload-time = "2025-10-14T04:41:51.051Z" }, + { url = "https://files.pythonhosted.org/packages/da/5f/6b8f83a55bb8278772c5ae54a577f3099025f9ade59d0136ac24a0df4bde/charset_normalizer-3.4.4-cp314-cp314-win_arm64.whl", hash = "sha256:de00632ca48df9daf77a2c65a484531649261ec9f25489917f09e455cb09ddb2", size = 100743, upload-time = "2025-10-14T04:41:52.122Z" }, + { url = "https://files.pythonhosted.org/packages/0a/4c/925909008ed5a988ccbb72dcc897407e5d6d3bd72410d69e051fc0c14647/charset_normalizer-3.4.4-py3-none-any.whl", hash = "sha256:7a32c560861a02ff789ad905a2fe94e3f840803362c84fecf1851cb4cf3dc37f", size = 53402, upload-time = "2025-10-14T04:42:31.76Z" }, +] + +[[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" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +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 = "contourpy" +version = "1.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/35/0167aad910bbdb9599272bd96d01a9ec6852f36b9455cf2ca67bd4cc2d23/contourpy-1.3.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:177fb367556747a686509d6fef71d221a4b198a3905fe824430e5ea0fda54eb5", size = 293257, upload-time = "2025-07-26T12:01:39.367Z" }, + { url = "https://files.pythonhosted.org/packages/96/e4/7adcd9c8362745b2210728f209bfbcf7d91ba868a2c5f40d8b58f54c509b/contourpy-1.3.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d002b6f00d73d69333dac9d0b8d5e84d9724ff9ef044fd63c5986e62b7c9e1b1", size = 274034, upload-time = "2025-07-26T12:01:40.645Z" }, + { url = "https://files.pythonhosted.org/packages/73/23/90e31ceeed1de63058a02cb04b12f2de4b40e3bef5e082a7c18d9c8ae281/contourpy-1.3.3-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:348ac1f5d4f1d66d3322420f01d42e43122f43616e0f194fc1c9f5d830c5b286", size = 334672, upload-time = "2025-07-26T12:01:41.942Z" }, + { url = "https://files.pythonhosted.org/packages/ed/93/b43d8acbe67392e659e1d984700e79eb67e2acb2bd7f62012b583a7f1b55/contourpy-1.3.3-cp313-cp313-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:655456777ff65c2c548b7c454af9c6f33f16c8884f11083244b5819cc214f1b5", size = 381234, upload-time = "2025-07-26T12:01:43.499Z" }, + { url = "https://files.pythonhosted.org/packages/46/3b/bec82a3ea06f66711520f75a40c8fc0b113b2a75edb36aa633eb11c4f50f/contourpy-1.3.3-cp313-cp313-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:644a6853d15b2512d67881586bd03f462c7ab755db95f16f14d7e238f2852c67", size = 385169, upload-time = "2025-07-26T12:01:45.219Z" }, + { url = "https://files.pythonhosted.org/packages/4b/32/e0f13a1c5b0f8572d0ec6ae2f6c677b7991fafd95da523159c19eff0696a/contourpy-1.3.3-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4debd64f124ca62069f313a9cb86656ff087786016d76927ae2cf37846b006c9", size = 362859, upload-time = "2025-07-26T12:01:46.519Z" }, + { url = "https://files.pythonhosted.org/packages/33/71/e2a7945b7de4e58af42d708a219f3b2f4cff7386e6b6ab0a0fa0033c49a9/contourpy-1.3.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a15459b0f4615b00bbd1e91f1b9e19b7e63aea7483d03d804186f278c0af2659", size = 1332062, upload-time = "2025-07-26T12:01:48.964Z" }, + { url = "https://files.pythonhosted.org/packages/12/fc/4e87ac754220ccc0e807284f88e943d6d43b43843614f0a8afa469801db0/contourpy-1.3.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ca0fdcd73925568ca027e0b17ab07aad764be4706d0a925b89227e447d9737b7", size = 1403932, upload-time = "2025-07-26T12:01:51.979Z" }, + { url = "https://files.pythonhosted.org/packages/a6/2e/adc197a37443f934594112222ac1aa7dc9a98faf9c3842884df9a9d8751d/contourpy-1.3.3-cp313-cp313-win32.whl", hash = "sha256:b20c7c9a3bf701366556e1b1984ed2d0cedf999903c51311417cf5f591d8c78d", size = 185024, upload-time = "2025-07-26T12:01:53.245Z" }, + { url = "https://files.pythonhosted.org/packages/18/0b/0098c214843213759692cc638fce7de5c289200a830e5035d1791d7a2338/contourpy-1.3.3-cp313-cp313-win_amd64.whl", hash = "sha256:1cadd8b8969f060ba45ed7c1b714fe69185812ab43bd6b86a9123fe8f99c3263", size = 226578, upload-time = "2025-07-26T12:01:54.422Z" }, + { url = "https://files.pythonhosted.org/packages/8a/9a/2f6024a0c5995243cd63afdeb3651c984f0d2bc727fd98066d40e141ad73/contourpy-1.3.3-cp313-cp313-win_arm64.whl", hash = "sha256:fd914713266421b7536de2bfa8181aa8c699432b6763a0ea64195ebe28bff6a9", size = 193524, upload-time = "2025-07-26T12:01:55.73Z" }, + { url = "https://files.pythonhosted.org/packages/c0/b3/f8a1a86bd3298513f500e5b1f5fd92b69896449f6cab6a146a5d52715479/contourpy-1.3.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:88df9880d507169449d434c293467418b9f6cbe82edd19284aa0409e7fdb933d", size = 306730, upload-time = "2025-07-26T12:01:57.051Z" }, + { url = "https://files.pythonhosted.org/packages/3f/11/4780db94ae62fc0c2053909b65dc3246bd7cecfc4f8a20d957ad43aa4ad8/contourpy-1.3.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:d06bb1f751ba5d417047db62bca3c8fde202b8c11fb50742ab3ab962c81e8216", size = 287897, upload-time = "2025-07-26T12:01:58.663Z" }, + { url = "https://files.pythonhosted.org/packages/ae/15/e59f5f3ffdd6f3d4daa3e47114c53daabcb18574a26c21f03dc9e4e42ff0/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e4e6b05a45525357e382909a4c1600444e2a45b4795163d3b22669285591c1ae", size = 326751, upload-time = "2025-07-26T12:02:00.343Z" }, + { url = "https://files.pythonhosted.org/packages/0f/81/03b45cfad088e4770b1dcf72ea78d3802d04200009fb364d18a493857210/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ab3074b48c4e2cf1a960e6bbeb7f04566bf36b1861d5c9d4d8ac04b82e38ba20", size = 375486, upload-time = "2025-07-26T12:02:02.128Z" }, + { url = "https://files.pythonhosted.org/packages/0c/ba/49923366492ffbdd4486e970d421b289a670ae8cf539c1ea9a09822b371a/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6c3d53c796f8647d6deb1abe867daeb66dcc8a97e8455efa729516b997b8ed99", size = 388106, upload-time = "2025-07-26T12:02:03.615Z" }, + { url = "https://files.pythonhosted.org/packages/9f/52/5b00ea89525f8f143651f9f03a0df371d3cbd2fccd21ca9b768c7a6500c2/contourpy-1.3.3-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:50ed930df7289ff2a8d7afeb9603f8289e5704755c7e5c3bbd929c90c817164b", size = 352548, upload-time = "2025-07-26T12:02:05.165Z" }, + { url = "https://files.pythonhosted.org/packages/32/1d/a209ec1a3a3452d490f6b14dd92e72280c99ae3d1e73da74f8277d4ee08f/contourpy-1.3.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:4feffb6537d64b84877da813a5c30f1422ea5739566abf0bd18065ac040e120a", size = 1322297, upload-time = "2025-07-26T12:02:07.379Z" }, + { url = "https://files.pythonhosted.org/packages/bc/9e/46f0e8ebdd884ca0e8877e46a3f4e633f6c9c8c4f3f6e72be3fe075994aa/contourpy-1.3.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:2b7e9480ffe2b0cd2e787e4df64270e3a0440d9db8dc823312e2c940c167df7e", size = 1391023, upload-time = "2025-07-26T12:02:10.171Z" }, + { url = "https://files.pythonhosted.org/packages/b9/70/f308384a3ae9cd2209e0849f33c913f658d3326900d0ff5d378d6a1422d2/contourpy-1.3.3-cp313-cp313t-win32.whl", hash = "sha256:283edd842a01e3dcd435b1c5116798d661378d83d36d337b8dde1d16a5fc9ba3", size = 196157, upload-time = "2025-07-26T12:02:11.488Z" }, + { url = "https://files.pythonhosted.org/packages/b2/dd/880f890a6663b84d9e34a6f88cded89d78f0091e0045a284427cb6b18521/contourpy-1.3.3-cp313-cp313t-win_amd64.whl", hash = "sha256:87acf5963fc2b34825e5b6b048f40e3635dd547f590b04d2ab317c2619ef7ae8", size = 240570, upload-time = "2025-07-26T12:02:12.754Z" }, + { url = "https://files.pythonhosted.org/packages/80/99/2adc7d8ffead633234817ef8e9a87115c8a11927a94478f6bb3d3f4d4f7d/contourpy-1.3.3-cp313-cp313t-win_arm64.whl", hash = "sha256:3c30273eb2a55024ff31ba7d052dde990d7d8e5450f4bbb6e913558b3d6c2301", size = 199713, upload-time = "2025-07-26T12:02:14.4Z" }, + { url = "https://files.pythonhosted.org/packages/72/8b/4546f3ab60f78c514ffb7d01a0bd743f90de36f0019d1be84d0a708a580a/contourpy-1.3.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:fde6c716d51c04b1c25d0b90364d0be954624a0ee9d60e23e850e8d48353d07a", size = 292189, upload-time = "2025-07-26T12:02:16.095Z" }, + { url = "https://files.pythonhosted.org/packages/fd/e1/3542a9cb596cadd76fcef413f19c79216e002623158befe6daa03dbfa88c/contourpy-1.3.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:cbedb772ed74ff5be440fa8eee9bd49f64f6e3fc09436d9c7d8f1c287b121d77", size = 273251, upload-time = "2025-07-26T12:02:17.524Z" }, + { url = "https://files.pythonhosted.org/packages/b1/71/f93e1e9471d189f79d0ce2497007731c1e6bf9ef6d1d61b911430c3db4e5/contourpy-1.3.3-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:22e9b1bd7a9b1d652cd77388465dc358dafcd2e217d35552424aa4f996f524f5", size = 335810, upload-time = "2025-07-26T12:02:18.9Z" }, + { url = "https://files.pythonhosted.org/packages/91/f9/e35f4c1c93f9275d4e38681a80506b5510e9327350c51f8d4a5a724d178c/contourpy-1.3.3-cp314-cp314-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a22738912262aa3e254e4f3cb079a95a67132fc5a063890e224393596902f5a4", size = 382871, upload-time = "2025-07-26T12:02:20.418Z" }, + { url = "https://files.pythonhosted.org/packages/b5/71/47b512f936f66a0a900d81c396a7e60d73419868fba959c61efed7a8ab46/contourpy-1.3.3-cp314-cp314-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:afe5a512f31ee6bd7d0dda52ec9864c984ca3d66664444f2d72e0dc4eb832e36", size = 386264, upload-time = "2025-07-26T12:02:21.916Z" }, + { url = "https://files.pythonhosted.org/packages/04/5f/9ff93450ba96b09c7c2b3f81c94de31c89f92292f1380261bd7195bea4ea/contourpy-1.3.3-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f64836de09927cba6f79dcd00fdd7d5329f3fccc633468507079c829ca4db4e3", size = 363819, upload-time = "2025-07-26T12:02:23.759Z" }, + { url = "https://files.pythonhosted.org/packages/3e/a6/0b185d4cc480ee494945cde102cb0149ae830b5fa17bf855b95f2e70ad13/contourpy-1.3.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:1fd43c3be4c8e5fd6e4f2baeae35ae18176cf2e5cced681cca908addf1cdd53b", size = 1333650, upload-time = "2025-07-26T12:02:26.181Z" }, + { url = "https://files.pythonhosted.org/packages/43/d7/afdc95580ca56f30fbcd3060250f66cedbde69b4547028863abd8aa3b47e/contourpy-1.3.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:6afc576f7b33cf00996e5c1102dc2a8f7cc89e39c0b55df93a0b78c1bd992b36", size = 1404833, upload-time = "2025-07-26T12:02:28.782Z" }, + { url = "https://files.pythonhosted.org/packages/e2/e2/366af18a6d386f41132a48f033cbd2102e9b0cf6345d35ff0826cd984566/contourpy-1.3.3-cp314-cp314-win32.whl", hash = "sha256:66c8a43a4f7b8df8b71ee1840e4211a3c8d93b214b213f590e18a1beca458f7d", size = 189692, upload-time = "2025-07-26T12:02:30.128Z" }, + { url = "https://files.pythonhosted.org/packages/7d/c2/57f54b03d0f22d4044b8afb9ca0e184f8b1afd57b4f735c2fa70883dc601/contourpy-1.3.3-cp314-cp314-win_amd64.whl", hash = "sha256:cf9022ef053f2694e31d630feaacb21ea24224be1c3ad0520b13d844274614fd", size = 232424, upload-time = "2025-07-26T12:02:31.395Z" }, + { url = "https://files.pythonhosted.org/packages/18/79/a9416650df9b525737ab521aa181ccc42d56016d2123ddcb7b58e926a42c/contourpy-1.3.3-cp314-cp314-win_arm64.whl", hash = "sha256:95b181891b4c71de4bb404c6621e7e2390745f887f2a026b2d99e92c17892339", size = 198300, upload-time = "2025-07-26T12:02:32.956Z" }, + { url = "https://files.pythonhosted.org/packages/1f/42/38c159a7d0f2b7b9c04c64ab317042bb6952b713ba875c1681529a2932fe/contourpy-1.3.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:33c82d0138c0a062380332c861387650c82e4cf1747aaa6938b9b6516762e772", size = 306769, upload-time = "2025-07-26T12:02:34.2Z" }, + { url = "https://files.pythonhosted.org/packages/c3/6c/26a8205f24bca10974e77460de68d3d7c63e282e23782f1239f226fcae6f/contourpy-1.3.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ea37e7b45949df430fe649e5de8351c423430046a2af20b1c1961cae3afcda77", size = 287892, upload-time = "2025-07-26T12:02:35.807Z" }, + { url = "https://files.pythonhosted.org/packages/66/06/8a475c8ab718ebfd7925661747dbb3c3ee9c82ac834ccb3570be49d129f4/contourpy-1.3.3-cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d304906ecc71672e9c89e87c4675dc5c2645e1f4269a5063b99b0bb29f232d13", size = 326748, upload-time = "2025-07-26T12:02:37.193Z" }, + { url = "https://files.pythonhosted.org/packages/b4/a3/c5ca9f010a44c223f098fccd8b158bb1cb287378a31ac141f04730dc49be/contourpy-1.3.3-cp314-cp314t-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ca658cd1a680a5c9ea96dc61cdbae1e85c8f25849843aa799dfd3cb370ad4fbe", size = 375554, upload-time = "2025-07-26T12:02:38.894Z" }, + { url = "https://files.pythonhosted.org/packages/80/5b/68bd33ae63fac658a4145088c1e894405e07584a316738710b636c6d0333/contourpy-1.3.3-cp314-cp314t-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ab2fd90904c503739a75b7c8c5c01160130ba67944a7b77bbf36ef8054576e7f", size = 388118, upload-time = "2025-07-26T12:02:40.642Z" }, + { url = "https://files.pythonhosted.org/packages/40/52/4c285a6435940ae25d7410a6c36bda5145839bc3f0beb20c707cda18b9d2/contourpy-1.3.3-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b7301b89040075c30e5768810bc96a8e8d78085b47d8be6e4c3f5a0b4ed478a0", size = 352555, upload-time = "2025-07-26T12:02:42.25Z" }, + { url = "https://files.pythonhosted.org/packages/24/ee/3e81e1dd174f5c7fefe50e85d0892de05ca4e26ef1c9a59c2a57e43b865a/contourpy-1.3.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:2a2a8b627d5cc6b7c41a4beff6c5ad5eb848c88255fda4a8745f7e901b32d8e4", size = 1322295, upload-time = "2025-07-26T12:02:44.668Z" }, + { url = "https://files.pythonhosted.org/packages/3c/b2/6d913d4d04e14379de429057cd169e5e00f6c2af3bb13e1710bcbdb5da12/contourpy-1.3.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:fd6ec6be509c787f1caf6b247f0b1ca598bef13f4ddeaa126b7658215529ba0f", size = 1391027, upload-time = "2025-07-26T12:02:47.09Z" }, + { url = "https://files.pythonhosted.org/packages/93/8a/68a4ec5c55a2971213d29a9374913f7e9f18581945a7a31d1a39b5d2dfe5/contourpy-1.3.3-cp314-cp314t-win32.whl", hash = "sha256:e74a9a0f5e3fff48fb5a7f2fd2b9b70a3fe014a67522f79b7cca4c0c7e43c9ae", size = 202428, upload-time = "2025-07-26T12:02:48.691Z" }, + { url = "https://files.pythonhosted.org/packages/fa/96/fd9f641ffedc4fa3ace923af73b9d07e869496c9cc7a459103e6e978992f/contourpy-1.3.3-cp314-cp314t-win_amd64.whl", hash = "sha256:13b68d6a62db8eafaebb8039218921399baf6e47bf85006fd8529f2a08ef33fc", size = 250331, upload-time = "2025-07-26T12:02:50.137Z" }, + { url = "https://files.pythonhosted.org/packages/ae/8c/469afb6465b853afff216f9528ffda78a915ff880ed58813ba4faf4ba0b6/contourpy-1.3.3-cp314-cp314t-win_arm64.whl", hash = "sha256:b7448cb5a725bb1e35ce88771b86fba35ef418952474492cf7c764059933ff8b", size = 203831, upload-time = "2025-07-26T12:02:51.449Z" }, +] + +[[package]] +name = "cycler" +version = "0.12.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/95/a3dbbb5028f35eafb79008e7522a75244477d2838f38cbb722248dabc2a8/cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c", size = 7615, upload-time = "2023-10-07T05:32:18.335Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, +] + +[[package]] +name = "dataclasses-json" +version = "0.6.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "marshmallow" }, + { name = "typing-inspect" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/64/a4/f71d9cf3a5ac257c993b5ca3f93df5f7fb395c725e7f1e6479d2514173c3/dataclasses_json-0.6.7.tar.gz", hash = "sha256:b6b3e528266ea45b9535223bc53ca645f5208833c29229e847b3f26a1cc55fc0", size = 32227, upload-time = "2024-06-09T16:20:19.103Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/be/d0d44e092656fe7a06b55e6103cbce807cdbdee17884a5367c68c9860853/dataclasses_json-0.6.7-py3-none-any.whl", hash = "sha256:0dbf33f26c8d5305befd61b39d2b3414e8a407bedc2834dea9b8d642666fb40a", size = 28686, upload-time = "2024-06-09T16:20:16.715Z" }, +] + +[[package]] +name = "etils" +version = "1.13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9b/a0/522bbff0f3cdd37968f90dd7f26c7aa801ed87f5ba335f156de7f2b88a48/etils-1.13.0.tar.gz", hash = "sha256:a5b60c71f95bcd2d43d4e9fb3dc3879120c1f60472bb5ce19f7a860b1d44f607", size = 106368, upload-time = "2025-07-15T10:29:10.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/98/87b5946356095738cb90a6df7b35ff69ac5750f6e783d5fbcc5cb3b6cbd7/etils-1.13.0-py3-none-any.whl", hash = "sha256:d9cd4f40fbe77ad6613b7348a18132cc511237b6c076dbb89105c0b520a4c6bb", size = 170603, upload-time = "2025-07-15T10:29:09.076Z" }, +] + +[package.optional-dependencies] +epath = [ + { name = "fsspec" }, + { name = "importlib-resources" }, + { name = "typing-extensions" }, + { name = "zipp" }, +] +epy = [ + { name = "typing-extensions" }, +] + +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, +] + +[[package]] +name = "flax" +version = "0.12.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax" }, + { name = "msgpack" }, + { name = "numpy" }, + { name = "optax" }, + { name = "orbax-checkpoint" }, + { name = "orbax-export" }, + { name = "pyyaml" }, + { name = "rich" }, + { name = "tensorstore" }, + { name = "treescope" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4d/ae/6cbdc663dd3b66dfb0f556342813639d8d04d2fdb7660dbfa75e4000ebf1/flax-0.12.3.tar.gz", hash = "sha256:8b9ca031d0a4759b0719718de578b27703a8b2d6105552b30f98989ce1be04f0", size = 5016605, upload-time = "2026-01-27T21:03:51.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4a/13/b3d5c22d8d4f266eba54795130afda106b12cce0fb92f3e196984725e024/flax-0.12.3-py3-none-any.whl", hash = "sha256:add0432489da2a532357d8a3a014a349d70774156a5a0a594706ba4319efa599", size = 490531, upload-time = "2026-01-27T21:03:50.295Z" }, +] + +[[package]] +name = "fonttools" +version = "4.61.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ec/ca/cf17b88a8df95691275a3d77dc0a5ad9907f328ae53acbe6795da1b2f5ed/fonttools-4.61.1.tar.gz", hash = "sha256:6675329885c44657f826ef01d9e4fb33b9158e9d93c537d84ad8399539bc6f69", size = 3565756, upload-time = "2025-12-12T17:31:24.246Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/cf/00ba28b0990982530addb8dc3e9e6f2fa9cb5c20df2abdda7baa755e8fe1/fonttools-4.61.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c56c488ab471628ff3bfa80964372fc13504ece601e0d97a78ee74126b2045c", size = 2846454, upload-time = "2025-12-12T17:30:24.938Z" }, + { url = "https://files.pythonhosted.org/packages/5a/ca/468c9a8446a2103ae645d14fee3f610567b7042aba85031c1c65e3ef7471/fonttools-4.61.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:dc492779501fa723b04d0ab1f5be046797fee17d27700476edc7ee9ae535a61e", size = 2398191, upload-time = "2025-12-12T17:30:27.343Z" }, + { url = "https://files.pythonhosted.org/packages/a3/4b/d67eedaed19def5967fade3297fed8161b25ba94699efc124b14fb68cdbc/fonttools-4.61.1-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:64102ca87e84261419c3747a0d20f396eb024bdbeb04c2bfb37e2891f5fadcb5", size = 4928410, upload-time = "2025-12-12T17:30:29.771Z" }, + { url = "https://files.pythonhosted.org/packages/b0/8d/6fb3494dfe61a46258cd93d979cf4725ded4eb46c2a4ca35e4490d84daea/fonttools-4.61.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4c1b526c8d3f615a7b1867f38a9410849c8f4aef078535742198e942fba0e9bd", size = 4984460, upload-time = "2025-12-12T17:30:32.073Z" }, + { url = "https://files.pythonhosted.org/packages/f7/f1/a47f1d30b3dc00d75e7af762652d4cbc3dff5c2697a0dbd5203c81afd9c3/fonttools-4.61.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:41ed4b5ec103bd306bb68f81dc166e77409e5209443e5773cb4ed837bcc9b0d3", size = 4925800, upload-time = "2025-12-12T17:30:34.339Z" }, + { url = "https://files.pythonhosted.org/packages/a7/01/e6ae64a0981076e8a66906fab01539799546181e32a37a0257b77e4aa88b/fonttools-4.61.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b501c862d4901792adaec7c25b1ecc749e2662543f68bb194c42ba18d6eec98d", size = 5067859, upload-time = "2025-12-12T17:30:36.593Z" }, + { url = "https://files.pythonhosted.org/packages/73/aa/28e40b8d6809a9b5075350a86779163f074d2b617c15d22343fce81918db/fonttools-4.61.1-cp313-cp313-win32.whl", hash = "sha256:4d7092bb38c53bbc78e9255a59158b150bcdc115a1e3b3ce0b5f267dc35dd63c", size = 2267821, upload-time = "2025-12-12T17:30:38.478Z" }, + { url = "https://files.pythonhosted.org/packages/1a/59/453c06d1d83dc0951b69ef692d6b9f1846680342927df54e9a1ca91c6f90/fonttools-4.61.1-cp313-cp313-win_amd64.whl", hash = "sha256:21e7c8d76f62ab13c9472ccf74515ca5b9a761d1bde3265152a6dc58700d895b", size = 2318169, upload-time = "2025-12-12T17:30:40.951Z" }, + { url = "https://files.pythonhosted.org/packages/32/8f/4e7bf82c0cbb738d3c2206c920ca34ca74ef9dabde779030145d28665104/fonttools-4.61.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:fff4f534200a04b4a36e7ae3cb74493afe807b517a09e99cb4faa89a34ed6ecd", size = 2846094, upload-time = "2025-12-12T17:30:43.511Z" }, + { url = "https://files.pythonhosted.org/packages/71/09/d44e45d0a4f3a651f23a1e9d42de43bc643cce2971b19e784cc67d823676/fonttools-4.61.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:d9203500f7c63545b4ce3799319fe4d9feb1a1b89b28d3cb5abd11b9dd64147e", size = 2396589, upload-time = "2025-12-12T17:30:45.681Z" }, + { url = "https://files.pythonhosted.org/packages/89/18/58c64cafcf8eb677a99ef593121f719e6dcbdb7d1c594ae5a10d4997ca8a/fonttools-4.61.1-cp314-cp314-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:fa646ecec9528bef693415c79a86e733c70a4965dd938e9a226b0fc64c9d2e6c", size = 4877892, upload-time = "2025-12-12T17:30:47.709Z" }, + { url = "https://files.pythonhosted.org/packages/8a/ec/9e6b38c7ba1e09eb51db849d5450f4c05b7e78481f662c3b79dbde6f3d04/fonttools-4.61.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:11f35ad7805edba3aac1a3710d104592df59f4b957e30108ae0ba6c10b11dd75", size = 4972884, upload-time = "2025-12-12T17:30:49.656Z" }, + { url = "https://files.pythonhosted.org/packages/5e/87/b5339da8e0256734ba0dbbf5b6cdebb1dd79b01dc8c270989b7bcd465541/fonttools-4.61.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b931ae8f62db78861b0ff1ac017851764602288575d65b8e8ff1963fed419063", size = 4924405, upload-time = "2025-12-12T17:30:51.735Z" }, + { url = "https://files.pythonhosted.org/packages/0b/47/e3409f1e1e69c073a3a6fd8cb886eb18c0bae0ee13db2c8d5e7f8495e8b7/fonttools-4.61.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:b148b56f5de675ee16d45e769e69f87623a4944f7443850bf9a9376e628a89d2", size = 5035553, upload-time = "2025-12-12T17:30:54.823Z" }, + { url = "https://files.pythonhosted.org/packages/bf/b6/1f6600161b1073a984294c6c031e1a56ebf95b6164249eecf30012bb2e38/fonttools-4.61.1-cp314-cp314-win32.whl", hash = "sha256:9b666a475a65f4e839d3d10473fad6d47e0a9db14a2f4a224029c5bfde58ad2c", size = 2271915, upload-time = "2025-12-12T17:30:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/52/7b/91e7b01e37cc8eb0e1f770d08305b3655e4f002fc160fb82b3390eabacf5/fonttools-4.61.1-cp314-cp314-win_amd64.whl", hash = "sha256:4f5686e1fe5fce75d82d93c47a438a25bf0d1319d2843a926f741140b2b16e0c", size = 2323487, upload-time = "2025-12-12T17:30:59.804Z" }, + { url = "https://files.pythonhosted.org/packages/39/5c/908ad78e46c61c3e3ed70c3b58ff82ab48437faf84ec84f109592cabbd9f/fonttools-4.61.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:e76ce097e3c57c4bcb67c5aa24a0ecdbd9f74ea9219997a707a4061fbe2707aa", size = 2929571, upload-time = "2025-12-12T17:31:02.574Z" }, + { url = "https://files.pythonhosted.org/packages/bd/41/975804132c6dea64cdbfbaa59f3518a21c137a10cccf962805b301ac6ab2/fonttools-4.61.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:9cfef3ab326780c04d6646f68d4b4742aae222e8b8ea1d627c74e38afcbc9d91", size = 2435317, upload-time = "2025-12-12T17:31:04.974Z" }, + { url = "https://files.pythonhosted.org/packages/b0/5a/aef2a0a8daf1ebaae4cfd83f84186d4a72ee08fd6a8451289fcd03ffa8a4/fonttools-4.61.1-cp314-cp314t-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:a75c301f96db737e1c5ed5fd7d77d9c34466de16095a266509e13da09751bd19", size = 4882124, upload-time = "2025-12-12T17:31:07.456Z" }, + { url = "https://files.pythonhosted.org/packages/80/33/d6db3485b645b81cea538c9d1c9219d5805f0877fda18777add4671c5240/fonttools-4.61.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:91669ccac46bbc1d09e9273546181919064e8df73488ea087dcac3e2968df9ba", size = 5100391, upload-time = "2025-12-12T17:31:09.732Z" }, + { url = "https://files.pythonhosted.org/packages/6c/d6/675ba631454043c75fcf76f0ca5463eac8eb0666ea1d7badae5fea001155/fonttools-4.61.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:c33ab3ca9d3ccd581d58e989d67554e42d8d4ded94ab3ade3508455fe70e65f7", size = 4978800, upload-time = "2025-12-12T17:31:11.681Z" }, + { url = "https://files.pythonhosted.org/packages/7f/33/d3ec753d547a8d2bdaedd390d4a814e8d5b45a093d558f025c6b990b554c/fonttools-4.61.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:664c5a68ec406f6b1547946683008576ef8b38275608e1cee6c061828171c118", size = 5006426, upload-time = "2025-12-12T17:31:13.764Z" }, + { url = "https://files.pythonhosted.org/packages/b4/40/cc11f378b561a67bea850ab50063366a0d1dd3f6d0a30ce0f874b0ad5664/fonttools-4.61.1-cp314-cp314t-win32.whl", hash = "sha256:aed04cabe26f30c1647ef0e8fbb207516fd40fe9472e9439695f5c6998e60ac5", size = 2335377, upload-time = "2025-12-12T17:31:16.49Z" }, + { url = "https://files.pythonhosted.org/packages/e4/ff/c9a2b66b39f8628531ea58b320d66d951267c98c6a38684daa8f50fb02f8/fonttools-4.61.1-cp314-cp314t-win_amd64.whl", hash = "sha256:2180f14c141d2f0f3da43f3a81bc8aa4684860f6b0e6f9e165a4831f24e6a23b", size = 2400613, upload-time = "2025-12-12T17:31:18.769Z" }, + { url = "https://files.pythonhosted.org/packages/c7/4e/ce75a57ff3aebf6fc1f4e9d508b8e5810618a33d900ad6c19eb30b290b97/fonttools-4.61.1-py3-none-any.whl", hash = "sha256:17d2bf5d541add43822bcf0c43d7d847b160c9bb01d15d5007d84e2217aaa371", size = 1148996, upload-time = "2025-12-12T17:31:21.03Z" }, +] + +[[package]] +name = "frozenlist" +version = "1.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2d/f5/c831fac6cc817d26fd54c7eaccd04ef7e0288806943f7cc5bbf69f3ac1f0/frozenlist-1.8.0.tar.gz", hash = "sha256:3ede829ed8d842f6cd48fc7081d7a41001a56f1f38603f9d49bf3020d59a31ad", size = 45875, upload-time = "2025-10-06T05:38:17.865Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2d/40/0832c31a37d60f60ed79e9dfb5a92e1e2af4f40a16a29abcc7992af9edff/frozenlist-1.8.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8d92f1a84bb12d9e56f818b3a746f3efba93c1b63c8387a73dde655e1e42282a", size = 85717, upload-time = "2025-10-06T05:36:27.341Z" }, + { url = "https://files.pythonhosted.org/packages/30/ba/b0b3de23f40bc55a7057bd38434e25c34fa48e17f20ee273bbde5e0650f3/frozenlist-1.8.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:96153e77a591c8adc2ee805756c61f59fef4cf4073a9275ee86fe8cba41241f7", size = 49651, upload-time = "2025-10-06T05:36:28.855Z" }, + { url = "https://files.pythonhosted.org/packages/0c/ab/6e5080ee374f875296c4243c381bbdef97a9ac39c6e3ce1d5f7d42cb78d6/frozenlist-1.8.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f21f00a91358803399890ab167098c131ec2ddd5f8f5fd5fe9c9f2c6fcd91e40", size = 49417, upload-time = "2025-10-06T05:36:29.877Z" }, + { url = "https://files.pythonhosted.org/packages/d5/4e/e4691508f9477ce67da2015d8c00acd751e6287739123113a9fca6f1604e/frozenlist-1.8.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:fb30f9626572a76dfe4293c7194a09fb1fe93ba94c7d4f720dfae3b646b45027", size = 234391, upload-time = "2025-10-06T05:36:31.301Z" }, + { url = "https://files.pythonhosted.org/packages/40/76/c202df58e3acdf12969a7895fd6f3bc016c642e6726aa63bd3025e0fc71c/frozenlist-1.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eaa352d7047a31d87dafcacbabe89df0aa506abb5b1b85a2fb91bc3faa02d822", size = 233048, upload-time = "2025-10-06T05:36:32.531Z" }, + { url = "https://files.pythonhosted.org/packages/f9/c0/8746afb90f17b73ca5979c7a3958116e105ff796e718575175319b5bb4ce/frozenlist-1.8.0-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:03ae967b4e297f58f8c774c7eabcce57fe3c2434817d4385c50661845a058121", size = 226549, upload-time = "2025-10-06T05:36:33.706Z" }, + { url = "https://files.pythonhosted.org/packages/7e/eb/4c7eefc718ff72f9b6c4893291abaae5fbc0c82226a32dcd8ef4f7a5dbef/frozenlist-1.8.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f6292f1de555ffcc675941d65fffffb0a5bcd992905015f85d0592201793e0e5", size = 239833, upload-time = "2025-10-06T05:36:34.947Z" }, + { url = "https://files.pythonhosted.org/packages/c2/4e/e5c02187cf704224f8b21bee886f3d713ca379535f16893233b9d672ea71/frozenlist-1.8.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:29548f9b5b5e3460ce7378144c3010363d8035cea44bc0bf02d57f5a685e084e", size = 245363, upload-time = "2025-10-06T05:36:36.534Z" }, + { url = "https://files.pythonhosted.org/packages/1f/96/cb85ec608464472e82ad37a17f844889c36100eed57bea094518bf270692/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ec3cc8c5d4084591b4237c0a272cc4f50a5b03396a47d9caaf76f5d7b38a4f11", size = 229314, upload-time = "2025-10-06T05:36:38.582Z" }, + { url = "https://files.pythonhosted.org/packages/5d/6f/4ae69c550e4cee66b57887daeebe006fe985917c01d0fff9caab9883f6d0/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:517279f58009d0b1f2e7c1b130b377a349405da3f7621ed6bfae50b10adf20c1", size = 243365, upload-time = "2025-10-06T05:36:40.152Z" }, + { url = "https://files.pythonhosted.org/packages/7a/58/afd56de246cf11780a40a2c28dc7cbabbf06337cc8ddb1c780a2d97e88d8/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:db1e72ede2d0d7ccb213f218df6a078a9c09a7de257c2fe8fcef16d5925230b1", size = 237763, upload-time = "2025-10-06T05:36:41.355Z" }, + { url = "https://files.pythonhosted.org/packages/cb/36/cdfaf6ed42e2644740d4a10452d8e97fa1c062e2a8006e4b09f1b5fd7d63/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:b4dec9482a65c54a5044486847b8a66bf10c9cb4926d42927ec4e8fd5db7fed8", size = 240110, upload-time = "2025-10-06T05:36:42.716Z" }, + { url = "https://files.pythonhosted.org/packages/03/a8/9ea226fbefad669f11b52e864c55f0bd57d3c8d7eb07e9f2e9a0b39502e1/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:21900c48ae04d13d416f0e1e0c4d81f7931f73a9dfa0b7a8746fb2fe7dd970ed", size = 233717, upload-time = "2025-10-06T05:36:44.251Z" }, + { url = "https://files.pythonhosted.org/packages/1e/0b/1b5531611e83ba7d13ccc9988967ea1b51186af64c42b7a7af465dcc9568/frozenlist-1.8.0-cp313-cp313-win32.whl", hash = "sha256:8b7b94a067d1c504ee0b16def57ad5738701e4ba10cec90529f13fa03c833496", size = 39628, upload-time = "2025-10-06T05:36:45.423Z" }, + { url = "https://files.pythonhosted.org/packages/d8/cf/174c91dbc9cc49bc7b7aab74d8b734e974d1faa8f191c74af9b7e80848e6/frozenlist-1.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:878be833caa6a3821caf85eb39c5ba92d28e85df26d57afb06b35b2efd937231", size = 43882, upload-time = "2025-10-06T05:36:46.796Z" }, + { url = "https://files.pythonhosted.org/packages/c1/17/502cd212cbfa96eb1388614fe39a3fc9ab87dbbe042b66f97acb57474834/frozenlist-1.8.0-cp313-cp313-win_arm64.whl", hash = "sha256:44389d135b3ff43ba8cc89ff7f51f5a0bb6b63d829c8300f79a2fe4fe61bcc62", size = 39676, upload-time = "2025-10-06T05:36:47.8Z" }, + { url = "https://files.pythonhosted.org/packages/d2/5c/3bbfaa920dfab09e76946a5d2833a7cbdf7b9b4a91c714666ac4855b88b4/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:e25ac20a2ef37e91c1b39938b591457666a0fa835c7783c3a8f33ea42870db94", size = 89235, upload-time = "2025-10-06T05:36:48.78Z" }, + { url = "https://files.pythonhosted.org/packages/d2/d6/f03961ef72166cec1687e84e8925838442b615bd0b8854b54923ce5b7b8a/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:07cdca25a91a4386d2e76ad992916a85038a9b97561bf7a3fd12d5d9ce31870c", size = 50742, upload-time = "2025-10-06T05:36:49.837Z" }, + { url = "https://files.pythonhosted.org/packages/1e/bb/a6d12b7ba4c3337667d0e421f7181c82dda448ce4e7ad7ecd249a16fa806/frozenlist-1.8.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:4e0c11f2cc6717e0a741f84a527c52616140741cd812a50422f83dc31749fb52", size = 51725, upload-time = "2025-10-06T05:36:50.851Z" }, + { url = "https://files.pythonhosted.org/packages/bc/71/d1fed0ffe2c2ccd70b43714c6cab0f4188f09f8a67a7914a6b46ee30f274/frozenlist-1.8.0-cp313-cp313t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b3210649ee28062ea6099cfda39e147fa1bc039583c8ee4481cb7811e2448c51", size = 284533, upload-time = "2025-10-06T05:36:51.898Z" }, + { url = "https://files.pythonhosted.org/packages/c9/1f/fb1685a7b009d89f9bf78a42d94461bc06581f6e718c39344754a5d9bada/frozenlist-1.8.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:581ef5194c48035a7de2aefc72ac6539823bb71508189e5de01d60c9dcd5fa65", size = 292506, upload-time = "2025-10-06T05:36:53.101Z" }, + { url = "https://files.pythonhosted.org/packages/e6/3b/b991fe1612703f7e0d05c0cf734c1b77aaf7c7d321df4572e8d36e7048c8/frozenlist-1.8.0-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3ef2d026f16a2b1866e1d86fc4e1291e1ed8a387b2c333809419a2f8b3a77b82", size = 274161, upload-time = "2025-10-06T05:36:54.309Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ec/c5c618767bcdf66e88945ec0157d7f6c4a1322f1473392319b7a2501ded7/frozenlist-1.8.0-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5500ef82073f599ac84d888e3a8c1f77ac831183244bfd7f11eaa0289fb30714", size = 294676, upload-time = "2025-10-06T05:36:55.566Z" }, + { url = "https://files.pythonhosted.org/packages/7c/ce/3934758637d8f8a88d11f0585d6495ef54b2044ed6ec84492a91fa3b27aa/frozenlist-1.8.0-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:50066c3997d0091c411a66e710f4e11752251e6d2d73d70d8d5d4c76442a199d", size = 300638, upload-time = "2025-10-06T05:36:56.758Z" }, + { url = "https://files.pythonhosted.org/packages/fc/4f/a7e4d0d467298f42de4b41cbc7ddaf19d3cfeabaf9ff97c20c6c7ee409f9/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:5c1c8e78426e59b3f8005e9b19f6ff46e5845895adbde20ece9218319eca6506", size = 283067, upload-time = "2025-10-06T05:36:57.965Z" }, + { url = "https://files.pythonhosted.org/packages/dc/48/c7b163063d55a83772b268e6d1affb960771b0e203b632cfe09522d67ea5/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:eefdba20de0d938cec6a89bd4d70f346a03108a19b9df4248d3cf0d88f1b0f51", size = 292101, upload-time = "2025-10-06T05:36:59.237Z" }, + { url = "https://files.pythonhosted.org/packages/9f/d0/2366d3c4ecdc2fd391e0afa6e11500bfba0ea772764d631bbf82f0136c9d/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:cf253e0e1c3ceb4aaff6df637ce033ff6535fb8c70a764a8f46aafd3d6ab798e", size = 289901, upload-time = "2025-10-06T05:37:00.811Z" }, + { url = "https://files.pythonhosted.org/packages/b8/94/daff920e82c1b70e3618a2ac39fbc01ae3e2ff6124e80739ce5d71c9b920/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:032efa2674356903cd0261c4317a561a6850f3ac864a63fc1583147fb05a79b0", size = 289395, upload-time = "2025-10-06T05:37:02.115Z" }, + { url = "https://files.pythonhosted.org/packages/e3/20/bba307ab4235a09fdcd3cc5508dbabd17c4634a1af4b96e0f69bfe551ebd/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6da155091429aeba16851ecb10a9104a108bcd32f6c1642867eadaee401c1c41", size = 283659, upload-time = "2025-10-06T05:37:03.711Z" }, + { url = "https://files.pythonhosted.org/packages/fd/00/04ca1c3a7a124b6de4f8a9a17cc2fcad138b4608e7a3fc5877804b8715d7/frozenlist-1.8.0-cp313-cp313t-win32.whl", hash = "sha256:0f96534f8bfebc1a394209427d0f8a63d343c9779cda6fc25e8e121b5fd8555b", size = 43492, upload-time = "2025-10-06T05:37:04.915Z" }, + { url = "https://files.pythonhosted.org/packages/59/5e/c69f733a86a94ab10f68e496dc6b7e8bc078ebb415281d5698313e3af3a1/frozenlist-1.8.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5d63a068f978fc69421fb0e6eb91a9603187527c86b7cd3f534a5b77a592b888", size = 48034, upload-time = "2025-10-06T05:37:06.343Z" }, + { url = "https://files.pythonhosted.org/packages/16/6c/be9d79775d8abe79b05fa6d23da99ad6e7763a1d080fbae7290b286093fd/frozenlist-1.8.0-cp313-cp313t-win_arm64.whl", hash = "sha256:bf0a7e10b077bf5fb9380ad3ae8ce20ef919a6ad93b4552896419ac7e1d8e042", size = 41749, upload-time = "2025-10-06T05:37:07.431Z" }, + { url = "https://files.pythonhosted.org/packages/f1/c8/85da824b7e7b9b6e7f7705b2ecaf9591ba6f79c1177f324c2735e41d36a2/frozenlist-1.8.0-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:cee686f1f4cadeb2136007ddedd0aaf928ab95216e7691c63e50a8ec066336d0", size = 86127, upload-time = "2025-10-06T05:37:08.438Z" }, + { url = "https://files.pythonhosted.org/packages/8e/e8/a1185e236ec66c20afd72399522f142c3724c785789255202d27ae992818/frozenlist-1.8.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:119fb2a1bd47307e899c2fac7f28e85b9a543864df47aa7ec9d3c1b4545f096f", size = 49698, upload-time = "2025-10-06T05:37:09.48Z" }, + { url = "https://files.pythonhosted.org/packages/a1/93/72b1736d68f03fda5fdf0f2180fb6caaae3894f1b854d006ac61ecc727ee/frozenlist-1.8.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4970ece02dbc8c3a92fcc5228e36a3e933a01a999f7094ff7c23fbd2beeaa67c", size = 49749, upload-time = "2025-10-06T05:37:10.569Z" }, + { url = "https://files.pythonhosted.org/packages/a7/b2/fabede9fafd976b991e9f1b9c8c873ed86f202889b864756f240ce6dd855/frozenlist-1.8.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:cba69cb73723c3f329622e34bdbf5ce1f80c21c290ff04256cff1cd3c2036ed2", size = 231298, upload-time = "2025-10-06T05:37:11.993Z" }, + { url = "https://files.pythonhosted.org/packages/3a/3b/d9b1e0b0eed36e70477ffb8360c49c85c8ca8ef9700a4e6711f39a6e8b45/frozenlist-1.8.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:778a11b15673f6f1df23d9586f83c4846c471a8af693a22e066508b77d201ec8", size = 232015, upload-time = "2025-10-06T05:37:13.194Z" }, + { url = "https://files.pythonhosted.org/packages/dc/94/be719d2766c1138148564a3960fc2c06eb688da592bdc25adcf856101be7/frozenlist-1.8.0-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0325024fe97f94c41c08872db482cf8ac4800d80e79222c6b0b7b162d5b13686", size = 225038, upload-time = "2025-10-06T05:37:14.577Z" }, + { url = "https://files.pythonhosted.org/packages/e4/09/6712b6c5465f083f52f50cf74167b92d4ea2f50e46a9eea0523d658454ae/frozenlist-1.8.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:97260ff46b207a82a7567b581ab4190bd4dfa09f4db8a8b49d1a958f6aa4940e", size = 240130, upload-time = "2025-10-06T05:37:15.781Z" }, + { url = "https://files.pythonhosted.org/packages/f8/d4/cd065cdcf21550b54f3ce6a22e143ac9e4836ca42a0de1022da8498eac89/frozenlist-1.8.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:54b2077180eb7f83dd52c40b2750d0a9f175e06a42e3213ce047219de902717a", size = 242845, upload-time = "2025-10-06T05:37:17.037Z" }, + { url = "https://files.pythonhosted.org/packages/62/c3/f57a5c8c70cd1ead3d5d5f776f89d33110b1addae0ab010ad774d9a44fb9/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:2f05983daecab868a31e1da44462873306d3cbfd76d1f0b5b69c473d21dbb128", size = 229131, upload-time = "2025-10-06T05:37:18.221Z" }, + { url = "https://files.pythonhosted.org/packages/6c/52/232476fe9cb64f0742f3fde2b7d26c1dac18b6d62071c74d4ded55e0ef94/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:33f48f51a446114bc5d251fb2954ab0164d5be02ad3382abcbfe07e2531d650f", size = 240542, upload-time = "2025-10-06T05:37:19.771Z" }, + { url = "https://files.pythonhosted.org/packages/5f/85/07bf3f5d0fb5414aee5f47d33c6f5c77bfe49aac680bfece33d4fdf6a246/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:154e55ec0655291b5dd1b8731c637ecdb50975a2ae70c606d100750a540082f7", size = 237308, upload-time = "2025-10-06T05:37:20.969Z" }, + { url = "https://files.pythonhosted.org/packages/11/99/ae3a33d5befd41ac0ca2cc7fd3aa707c9c324de2e89db0e0f45db9a64c26/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:4314debad13beb564b708b4a496020e5306c7333fa9a3ab90374169a20ffab30", size = 238210, upload-time = "2025-10-06T05:37:22.252Z" }, + { url = "https://files.pythonhosted.org/packages/b2/60/b1d2da22f4970e7a155f0adde9b1435712ece01b3cd45ba63702aea33938/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:073f8bf8becba60aa931eb3bc420b217bb7d5b8f4750e6f8b3be7f3da85d38b7", size = 231972, upload-time = "2025-10-06T05:37:23.5Z" }, + { url = "https://files.pythonhosted.org/packages/3f/ab/945b2f32de889993b9c9133216c068b7fcf257d8595a0ac420ac8677cab0/frozenlist-1.8.0-cp314-cp314-win32.whl", hash = "sha256:bac9c42ba2ac65ddc115d930c78d24ab8d4f465fd3fc473cdedfccadb9429806", size = 40536, upload-time = "2025-10-06T05:37:25.581Z" }, + { url = "https://files.pythonhosted.org/packages/59/ad/9caa9b9c836d9ad6f067157a531ac48b7d36499f5036d4141ce78c230b1b/frozenlist-1.8.0-cp314-cp314-win_amd64.whl", hash = "sha256:3e0761f4d1a44f1d1a47996511752cf3dcec5bbdd9cc2b4fe595caf97754b7a0", size = 44330, upload-time = "2025-10-06T05:37:26.928Z" }, + { url = "https://files.pythonhosted.org/packages/82/13/e6950121764f2676f43534c555249f57030150260aee9dcf7d64efda11dd/frozenlist-1.8.0-cp314-cp314-win_arm64.whl", hash = "sha256:d1eaff1d00c7751b7c6662e9c5ba6eb2c17a2306ba5e2a37f24ddf3cc953402b", size = 40627, upload-time = "2025-10-06T05:37:28.075Z" }, + { url = "https://files.pythonhosted.org/packages/c0/c7/43200656ecc4e02d3f8bc248df68256cd9572b3f0017f0a0c4e93440ae23/frozenlist-1.8.0-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:d3bb933317c52d7ea5004a1c442eef86f426886fba134ef8cf4226ea6ee1821d", size = 89238, upload-time = "2025-10-06T05:37:29.373Z" }, + { url = "https://files.pythonhosted.org/packages/d1/29/55c5f0689b9c0fb765055629f472c0de484dcaf0acee2f7707266ae3583c/frozenlist-1.8.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:8009897cdef112072f93a0efdce29cd819e717fd2f649ee3016efd3cd885a7ed", size = 50738, upload-time = "2025-10-06T05:37:30.792Z" }, + { url = "https://files.pythonhosted.org/packages/ba/7d/b7282a445956506fa11da8c2db7d276adcbf2b17d8bb8407a47685263f90/frozenlist-1.8.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:2c5dcbbc55383e5883246d11fd179782a9d07a986c40f49abe89ddf865913930", size = 51739, upload-time = "2025-10-06T05:37:32.127Z" }, + { url = "https://files.pythonhosted.org/packages/62/1c/3d8622e60d0b767a5510d1d3cf21065b9db874696a51ea6d7a43180a259c/frozenlist-1.8.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:39ecbc32f1390387d2aa4f5a995e465e9e2f79ba3adcac92d68e3e0afae6657c", size = 284186, upload-time = "2025-10-06T05:37:33.21Z" }, + { url = "https://files.pythonhosted.org/packages/2d/14/aa36d5f85a89679a85a1d44cd7a6657e0b1c75f61e7cad987b203d2daca8/frozenlist-1.8.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:92db2bf818d5cc8d9c1f1fc56b897662e24ea5adb36ad1f1d82875bd64e03c24", size = 292196, upload-time = "2025-10-06T05:37:36.107Z" }, + { url = "https://files.pythonhosted.org/packages/05/23/6bde59eb55abd407d34f77d39a5126fb7b4f109a3f611d3929f14b700c66/frozenlist-1.8.0-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2dc43a022e555de94c3b68a4ef0b11c4f747d12c024a520c7101709a2144fb37", size = 273830, upload-time = "2025-10-06T05:37:37.663Z" }, + { url = "https://files.pythonhosted.org/packages/d2/3f/22cff331bfad7a8afa616289000ba793347fcd7bc275f3b28ecea2a27909/frozenlist-1.8.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cb89a7f2de3602cfed448095bab3f178399646ab7c61454315089787df07733a", size = 294289, upload-time = "2025-10-06T05:37:39.261Z" }, + { url = "https://files.pythonhosted.org/packages/a4/89/5b057c799de4838b6c69aa82b79705f2027615e01be996d2486a69ca99c4/frozenlist-1.8.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:33139dc858c580ea50e7e60a1b0ea003efa1fd42e6ec7fdbad78fff65fad2fd2", size = 300318, upload-time = "2025-10-06T05:37:43.213Z" }, + { url = "https://files.pythonhosted.org/packages/30/de/2c22ab3eb2a8af6d69dc799e48455813bab3690c760de58e1bf43b36da3e/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:168c0969a329b416119507ba30b9ea13688fafffac1b7822802537569a1cb0ef", size = 282814, upload-time = "2025-10-06T05:37:45.337Z" }, + { url = "https://files.pythonhosted.org/packages/59/f7/970141a6a8dbd7f556d94977858cfb36fa9b66e0892c6dd780d2219d8cd8/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:28bd570e8e189d7f7b001966435f9dac6718324b5be2990ac496cf1ea9ddb7fe", size = 291762, upload-time = "2025-10-06T05:37:46.657Z" }, + { url = "https://files.pythonhosted.org/packages/c1/15/ca1adae83a719f82df9116d66f5bb28bb95557b3951903d39135620ef157/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:b2a095d45c5d46e5e79ba1e5b9cb787f541a8dee0433836cea4b96a2c439dcd8", size = 289470, upload-time = "2025-10-06T05:37:47.946Z" }, + { url = "https://files.pythonhosted.org/packages/ac/83/dca6dc53bf657d371fbc88ddeb21b79891e747189c5de990b9dfff2ccba1/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:eab8145831a0d56ec9c4139b6c3e594c7a83c2c8be25d5bcf2d86136a532287a", size = 289042, upload-time = "2025-10-06T05:37:49.499Z" }, + { url = "https://files.pythonhosted.org/packages/96/52/abddd34ca99be142f354398700536c5bd315880ed0a213812bc491cff5e4/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:974b28cf63cc99dfb2188d8d222bc6843656188164848c4f679e63dae4b0708e", size = 283148, upload-time = "2025-10-06T05:37:50.745Z" }, + { url = "https://files.pythonhosted.org/packages/af/d3/76bd4ed4317e7119c2b7f57c3f6934aba26d277acc6309f873341640e21f/frozenlist-1.8.0-cp314-cp314t-win32.whl", hash = "sha256:342c97bf697ac5480c0a7ec73cd700ecfa5a8a40ac923bd035484616efecc2df", size = 44676, upload-time = "2025-10-06T05:37:52.222Z" }, + { url = "https://files.pythonhosted.org/packages/89/76/c615883b7b521ead2944bb3480398cbb07e12b7b4e4d073d3752eb721558/frozenlist-1.8.0-cp314-cp314t-win_amd64.whl", hash = "sha256:06be8f67f39c8b1dc671f5d83aaefd3358ae5cdcf8314552c57e7ed3e6475bdd", size = 49451, upload-time = "2025-10-06T05:37:53.425Z" }, + { url = "https://files.pythonhosted.org/packages/e0/a3/5982da14e113d07b325230f95060e2169f5311b1017ea8af2a29b374c289/frozenlist-1.8.0-cp314-cp314t-win_arm64.whl", hash = "sha256:102e6314ca4da683dca92e3b1355490fed5f313b768500084fbe6371fddfdb79", size = 42507, upload-time = "2025-10-06T05:37:54.513Z" }, + { url = "https://files.pythonhosted.org/packages/9a/9a/e35b4a917281c0b8419d4207f4334c8e8c5dbf4f3f5f9ada73958d937dcc/frozenlist-1.8.0-py3-none-any.whl", hash = "sha256:0c18a16eab41e82c295618a77502e17b195883241c563b00f0aa5106fc4eaa0d", size = 13409, upload-time = "2025-10-06T05:38:16.721Z" }, +] + +[[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 = "graphviz" +version = "0.21" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f8/b3/3ac91e9be6b761a4b30d66ff165e54439dcd48b83f4e20d644867215f6ca/graphviz-0.21.tar.gz", hash = "sha256:20743e7183be82aaaa8ad6c93f8893c923bd6658a04c32ee115edb3c8a835f78", size = 200434, upload-time = "2025-06-15T09:35:05.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/4c/e0ce1ef95d4000ebc1c11801f9b944fa5910ecc15b5e351865763d8657f8/graphviz-0.21-py3-none-any.whl", hash = "sha256:54f33de9f4f911d7e84e4191749cac8cc5653f815b06738c54db9a15ab8b1e42", size = 47300, upload-time = "2025-06-15T09:35:04.433Z" }, +] + +[[package]] +name = "grpcio" +version = "1.78.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/8a/3d098f35c143a89520e568e6539cc098fcd294495910e359889ce8741c84/grpcio-1.78.0.tar.gz", hash = "sha256:7382b95189546f375c174f53a5fa873cef91c4b8005faa05cc5b3beea9c4f1c5", size = 12852416, upload-time = "2026-02-06T09:57:18.093Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/05/a9/8f75894993895f361ed8636cd9237f4ab39ef87fd30db17467235ed1c045/grpcio-1.78.0-cp313-cp313-linux_armv7l.whl", hash = "sha256:ce3a90455492bf8bfa38e56fbbe1dbd4f872a3d8eeaf7337dc3b1c8aa28c271b", size = 5920143, upload-time = "2026-02-06T09:55:52.035Z" }, + { url = "https://files.pythonhosted.org/packages/55/06/0b78408e938ac424100100fd081189451b472236e8a3a1f6500390dc4954/grpcio-1.78.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:2bf5e2e163b356978b23652c4818ce4759d40f4712ee9ec5a83c4be6f8c23a3a", size = 11803926, upload-time = "2026-02-06T09:55:55.494Z" }, + { url = "https://files.pythonhosted.org/packages/88/93/b59fe7832ff6ae3c78b813ea43dac60e295fa03606d14d89d2e0ec29f4f3/grpcio-1.78.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8f2ac84905d12918e4e55a16da17939eb63e433dc11b677267c35568aa63fc84", size = 6478628, upload-time = "2026-02-06T09:55:58.533Z" }, + { url = "https://files.pythonhosted.org/packages/ed/df/e67e3734527f9926b7d9c0dde6cd998d1d26850c3ed8eeec81297967ac67/grpcio-1.78.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:b58f37edab4a3881bc6c9bca52670610e0c9ca14e2ea3cf9debf185b870457fb", size = 7173574, upload-time = "2026-02-06T09:56:01.786Z" }, + { url = "https://files.pythonhosted.org/packages/a6/62/cc03fffb07bfba982a9ec097b164e8835546980aec25ecfa5f9c1a47e022/grpcio-1.78.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:735e38e176a88ce41840c21bb49098ab66177c64c82426e24e0082500cc68af5", size = 6692639, upload-time = "2026-02-06T09:56:04.529Z" }, + { url = "https://files.pythonhosted.org/packages/bf/9a/289c32e301b85bdb67d7ec68b752155e674ee3ba2173a1858f118e399ef3/grpcio-1.78.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2045397e63a7a0ee7957c25f7dbb36ddc110e0cfb418403d110c0a7a68a844e9", size = 7268838, upload-time = "2026-02-06T09:56:08.397Z" }, + { url = "https://files.pythonhosted.org/packages/0e/79/1be93f32add280461fa4773880196572563e9c8510861ac2da0ea0f892b6/grpcio-1.78.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:a9f136fbafe7ccf4ac7e8e0c28b31066e810be52d6e344ef954a3a70234e1702", size = 8251878, upload-time = "2026-02-06T09:56:10.914Z" }, + { url = "https://files.pythonhosted.org/packages/65/65/793f8e95296ab92e4164593674ae6291b204bb5f67f9d4a711489cd30ffa/grpcio-1.78.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:748b6138585379c737adc08aeffd21222abbda1a86a0dca2a39682feb9196c20", size = 7695412, upload-time = "2026-02-06T09:56:13.593Z" }, + { url = "https://files.pythonhosted.org/packages/1c/9f/1e233fe697ecc82845942c2822ed06bb522e70d6771c28d5528e4c50f6a4/grpcio-1.78.0-cp313-cp313-win32.whl", hash = "sha256:271c73e6e5676afe4fc52907686670c7cea22ab2310b76a59b678403ed40d670", size = 4064899, upload-time = "2026-02-06T09:56:15.601Z" }, + { url = "https://files.pythonhosted.org/packages/4d/27/d86b89e36de8a951501fb06a0f38df19853210f341d0b28f83f4aa0ffa08/grpcio-1.78.0-cp313-cp313-win_amd64.whl", hash = "sha256:f2d4e43ee362adfc05994ed479334d5a451ab7bc3f3fee1b796b8ca66895acb4", size = 4797393, upload-time = "2026-02-06T09:56:17.882Z" }, + { url = "https://files.pythonhosted.org/packages/29/f2/b56e43e3c968bfe822fa6ce5bca10d5c723aa40875b48791ce1029bb78c7/grpcio-1.78.0-cp314-cp314-linux_armv7l.whl", hash = "sha256:e87cbc002b6f440482b3519e36e1313eb5443e9e9e73d6a52d43bd2004fcfd8e", size = 5920591, upload-time = "2026-02-06T09:56:20.758Z" }, + { url = "https://files.pythonhosted.org/packages/5d/81/1f3b65bd30c334167bfa8b0d23300a44e2725ce39bba5b76a2460d85f745/grpcio-1.78.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:c41bc64626db62e72afec66b0c8a0da76491510015417c127bfc53b2fe6d7f7f", size = 11813685, upload-time = "2026-02-06T09:56:24.315Z" }, + { url = "https://files.pythonhosted.org/packages/0e/1c/bbe2f8216a5bd3036119c544d63c2e592bdf4a8ec6e4a1867592f4586b26/grpcio-1.78.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8dfffba826efcf366b1e3ccc37e67afe676f290e13a3b48d31a46739f80a8724", size = 6487803, upload-time = "2026-02-06T09:56:27.367Z" }, + { url = "https://files.pythonhosted.org/packages/16/5c/a6b2419723ea7ddce6308259a55e8e7593d88464ce8db9f4aa857aba96fa/grpcio-1.78.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:74be1268d1439eaaf552c698cdb11cd594f0c49295ae6bb72c34ee31abbe611b", size = 7173206, upload-time = "2026-02-06T09:56:29.876Z" }, + { url = "https://files.pythonhosted.org/packages/df/1e/b8801345629a415ea7e26c83d75eb5dbe91b07ffe5210cc517348a8d4218/grpcio-1.78.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:be63c88b32e6c0f1429f1398ca5c09bc64b0d80950c8bb7807d7d7fb36fb84c7", size = 6693826, upload-time = "2026-02-06T09:56:32.305Z" }, + { url = "https://files.pythonhosted.org/packages/34/84/0de28eac0377742679a510784f049738a80424b17287739fc47d63c2439e/grpcio-1.78.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:3c586ac70e855c721bda8f548d38c3ca66ac791dc49b66a8281a1f99db85e452", size = 7277897, upload-time = "2026-02-06T09:56:34.915Z" }, + { url = "https://files.pythonhosted.org/packages/ca/9c/ad8685cfe20559a9edb66f735afdcb2b7d3de69b13666fdfc542e1916ebd/grpcio-1.78.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:35eb275bf1751d2ffbd8f57cdbc46058e857cf3971041521b78b7db94bdaf127", size = 8252404, upload-time = "2026-02-06T09:56:37.553Z" }, + { url = "https://files.pythonhosted.org/packages/3c/05/33a7a4985586f27e1de4803887c417ec7ced145ebd069bc38a9607059e2b/grpcio-1.78.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:207db540302c884b8848036b80db352a832b99dfdf41db1eb554c2c2c7800f65", size = 7696837, upload-time = "2026-02-06T09:56:40.173Z" }, + { url = "https://files.pythonhosted.org/packages/73/77/7382241caf88729b106e49e7d18e3116216c778e6a7e833826eb96de22f7/grpcio-1.78.0-cp314-cp314-win32.whl", hash = "sha256:57bab6deef2f4f1ca76cc04565df38dc5713ae6c17de690721bdf30cb1e0545c", size = 4142439, upload-time = "2026-02-06T09:56:43.258Z" }, + { url = "https://files.pythonhosted.org/packages/48/b2/b096ccce418882fbfda4f7496f9357aaa9a5af1896a9a7f60d9f2b275a06/grpcio-1.78.0-cp314-cp314-win_amd64.whl", hash = "sha256:dce09d6116df20a96acfdbf85e4866258c3758180e8c49845d6ba8248b6d0bbb", size = 4929852, upload-time = "2026-02-06T09:56:45.885Z" }, +] + +[[package]] +name = "grpcio-tools" +version = "1.78.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "grpcio" }, + { name = "protobuf" }, + { name = "setuptools" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8b/d1/cbefe328653f746fd319c4377836a25ba64226e41c6a1d7d5cdbc87a459f/grpcio_tools-1.78.0.tar.gz", hash = "sha256:4b0dd86560274316e155d925158276f8564508193088bc43e20d3f5dff956b2b", size = 5393026, upload-time = "2026-02-06T09:59:59.53Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/ce/17311fb77530420e2f441e916b347515133e83d21cd6cc77be04ce093d5b/grpcio_tools-1.78.0-cp313-cp313-linux_armv7l.whl", hash = "sha256:2d6de1cc23bdc1baafc23e201b1e48c617b8c1418b4d8e34cebf72141676e5fb", size = 2546284, upload-time = "2026-02-06T09:58:43.073Z" }, + { url = "https://files.pythonhosted.org/packages/1d/d3/79e101483115f0e78223397daef71751b75eba7e92a32060c10aae11ca64/grpcio_tools-1.78.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:2afeaad88040894c76656202ff832cb151bceb05c0e6907e539d129188b1e456", size = 5705653, upload-time = "2026-02-06T09:58:45.533Z" }, + { url = "https://files.pythonhosted.org/packages/8b/a7/52fa3ccb39ceeee6adc010056eadfbca8198651c113e418dafebbdf2b306/grpcio_tools-1.78.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:33cc593735c93c03d63efe7a8ba25f3c66f16c52f0651910712490244facad72", size = 2592788, upload-time = "2026-02-06T09:58:48.918Z" }, + { url = "https://files.pythonhosted.org/packages/68/08/682ff6bb548225513d73dc9403742d8975439d7469c673bc534b9bbc83a7/grpcio_tools-1.78.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:2921d7989c4d83b71f03130ab415fa4d66e6693b8b8a1fcbb7a1c67cff19b812", size = 2905157, upload-time = "2026-02-06T09:58:51.478Z" }, + { url = "https://files.pythonhosted.org/packages/b2/66/264f3836a96423b7018e5ada79d62576a6401f6da4e1f4975b18b2be1265/grpcio_tools-1.78.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e6a0df438e82c804c7b95e3f311c97c2f876dcc36376488d5b736b7bcf5a9b45", size = 2656166, upload-time = "2026-02-06T09:58:54.117Z" }, + { url = "https://files.pythonhosted.org/packages/f3/6b/f108276611522e03e98386b668cc7e575eff6952f2db9caa15b2a3b3e883/grpcio_tools-1.78.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e9c6070a9500798225191ef25d0055a15d2c01c9c8f2ee7b681fffa99c98c822", size = 3109110, upload-time = "2026-02-06T09:58:56.891Z" }, + { url = "https://files.pythonhosted.org/packages/6f/c7/cf048dbcd64b3396b3c860a2ffbcc67a8f8c87e736aaa74c2e505a7eee4c/grpcio_tools-1.78.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:394e8b57d85370a62e5b0a4d64c96fcf7568345c345d8590c821814d227ecf1d", size = 3657863, upload-time = "2026-02-06T09:58:59.176Z" }, + { url = "https://files.pythonhosted.org/packages/b6/37/e2736912c8fda57e2e57a66ea5e0bc8eb9a5fb7ded00e866ad22d50afb08/grpcio_tools-1.78.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:a3ef700293ab375e111a2909d87434ed0a0b086adf0ce67a8d9cf12ea7765e63", size = 3324748, upload-time = "2026-02-06T09:59:01.242Z" }, + { url = "https://files.pythonhosted.org/packages/1c/5d/726abc75bb5bfc2841e88ea05896e42f51ca7c30cb56da5c5b63058b3867/grpcio_tools-1.78.0-cp313-cp313-win32.whl", hash = "sha256:6993b960fec43a8d840ee5dc20247ef206c1a19587ea49fe5e6cc3d2a09c1585", size = 993074, upload-time = "2026-02-06T09:59:03.085Z" }, + { url = "https://files.pythonhosted.org/packages/c5/68/91b400bb360faf9b177ffb5540ec1c4d06ca923691ddf0f79e2c9683f4da/grpcio_tools-1.78.0-cp313-cp313-win_amd64.whl", hash = "sha256:275ce3c2978842a8cf9dd88dce954e836e590cf7029649ad5d1145b779039ed5", size = 1158185, upload-time = "2026-02-06T09:59:05.036Z" }, + { url = "https://files.pythonhosted.org/packages/cf/5e/278f3831c8d56bae02e3acc570465648eccf0a6bbedcb1733789ac966803/grpcio_tools-1.78.0-cp314-cp314-linux_armv7l.whl", hash = "sha256:8b080d0d072e6032708a3a91731b808074d7ab02ca8fb9847b6a011fdce64cd9", size = 2546270, upload-time = "2026-02-06T09:59:07.426Z" }, + { url = "https://files.pythonhosted.org/packages/a3/d9/68582f2952b914b60dddc18a2e3f9c6f09af9372b6f6120d6cf3ec7f8b4e/grpcio_tools-1.78.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:8c0ad8f8f133145cd7008b49cb611a5c6a9d89ab276c28afa17050516e801f79", size = 5705731, upload-time = "2026-02-06T09:59:09.856Z" }, + { url = "https://files.pythonhosted.org/packages/70/68/feb0f9a48818ee1df1e8b644069379a1e6ef5447b9b347c24e96fd258e5d/grpcio_tools-1.78.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2f8ea092a7de74c6359335d36f0674d939a3c7e1a550f4c2c9e80e0226de8fe4", size = 2593896, upload-time = "2026-02-06T09:59:12.23Z" }, + { url = "https://files.pythonhosted.org/packages/1f/08/a430d8d06e1b8d33f3e48d3f0cc28236723af2f35e37bd5c8db05df6c3aa/grpcio_tools-1.78.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:da422985e0cac822b41822f43429c19ecb27c81ffe3126d0b74e77edec452608", size = 2905298, upload-time = "2026-02-06T09:59:14.458Z" }, + { url = "https://files.pythonhosted.org/packages/71/0a/348c36a3eae101ca0c090c9c3bc96f2179adf59ee0c9262d11cdc7bfe7db/grpcio_tools-1.78.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4fab1faa3fbcb246263e68da7a8177d73772283f9db063fb8008517480888d26", size = 2656186, upload-time = "2026-02-06T09:59:16.949Z" }, + { url = "https://files.pythonhosted.org/packages/1d/3f/18219f331536fad4af6207ade04142292faa77b5cb4f4463787988963df8/grpcio_tools-1.78.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:dd9c094f73f734becae3f20f27d4944d3cd8fb68db7338ee6c58e62fc5c3d99f", size = 3109859, upload-time = "2026-02-06T09:59:19.202Z" }, + { url = "https://files.pythonhosted.org/packages/5b/d9/341ea20a44c8e5a3a18acc820b65014c2e3ea5b4f32a53d14864bcd236bc/grpcio_tools-1.78.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:2ed51ce6b833068f6c580b73193fc2ec16468e6bc18354bc2f83a58721195a58", size = 3657915, upload-time = "2026-02-06T09:59:21.839Z" }, + { url = "https://files.pythonhosted.org/packages/fb/f4/5978b0f91611a64371424c109dd0027b247e5b39260abad2eaee66b6aa37/grpcio_tools-1.78.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:05803a5cdafe77c8bdf36aa660ad7a6a1d9e49bc59ce45c1bade2a4698826599", size = 3324724, upload-time = "2026-02-06T09:59:24.402Z" }, + { url = "https://files.pythonhosted.org/packages/b2/80/96a324dba99cfbd20e291baf0b0ae719dbb62b76178c5ce6c788e7331cb1/grpcio_tools-1.78.0-cp314-cp314-win32.whl", hash = "sha256:f7c722e9ce6f11149ac5bddd5056e70aaccfd8168e74e9d34d8b8b588c3f5c7c", size = 1015505, upload-time = "2026-02-06T09:59:26.3Z" }, + { url = "https://files.pythonhosted.org/packages/3b/d1/909e6a05bfd44d46327dc4b8a78beb2bae4fb245ffab2772e350081aaf7e/grpcio_tools-1.78.0-cp314-cp314-win_amd64.whl", hash = "sha256:7d58ade518b546120ec8f0a8e006fc8076ae5df151250ebd7e82e9b5e152c229", size = 1190196, upload-time = "2026-02-06T09:59:28.359Z" }, +] + +[[package]] +name = "humanize" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/66/a3921783d54be8a6870ac4ccffcd15c4dc0dd7fcce51c6d63b8c63935276/humanize-4.15.0.tar.gz", hash = "sha256:1dd098483eb1c7ee8e32eb2e99ad1910baefa4b75c3aff3a82f4d78688993b10", size = 83599, upload-time = "2025-12-20T20:16:13.19Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c5/7b/bca5613a0c3b542420cf92bd5e5fb8ebd5435ce1011a091f66bb7693285e/humanize-4.15.0-py3-none-any.whl", hash = "sha256:b1186eb9f5a9749cd9cb8565aee77919dd7c8d076161cf44d70e59e3301e1769", size = 132203, upload-time = "2025-12-20T20:16:11.67Z" }, +] + +[[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 = "importlib-resources" +version = "6.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cf/8c/f834fbf984f691b4f7ff60f50b514cc3de5cc08abfc3295564dd89c5e2e7/importlib_resources-6.5.2.tar.gz", hash = "sha256:185f87adef5bcc288449d98fb4fba07cea78bc036455dd44c5fc4a2fe78fed2c", size = 44693, upload-time = "2025-01-03T18:51:56.698Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/ed/1f1afb2e9e7f38a545d628f864d562a5ae64fe6f7a10e28ffb9b185b4e89/importlib_resources-6.5.2-py3-none-any.whl", hash = "sha256:789cfdc3ed28c78b67a06acb8126751ced69a3d5f79c095a98298cd8a760ccec", size = 37461, upload-time = "2025-01-03T18:51:54.306Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +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 = "jax" +version = "0.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jaxlib" }, + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "opt-einsum" }, + { name = "scipy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/25/4d/f45853fdc2b811e78b866d5f80b8a21a848278361f66c066706132f415cf/jax-0.9.1.tar.gz", hash = "sha256:ce1b82477ee192f0b1d9801b095aa0cf3839bc1fe0cbc071c961a24b3ff30361", size = 2625994, upload-time = "2026-03-02T11:24:18.382Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/80/e4/88778c6a23b65224e5088e68fd0924e5bde2196a26e76edb3ea3543fed6a/jax-0.9.1-py3-none-any.whl", hash = "sha256:d11cb53d362912253013e8c4d6926cb9f3a4b59ab5b25a7dc08123567067d088", size = 3062162, upload-time = "2026-03-02T11:22:05.089Z" }, +] + +[package.optional-dependencies] +cuda12 = [ + { name = "jax-cuda12-plugin", extra = ["with-cuda"] }, + { name = "jaxlib" }, +] + +[[package]] +name = "jax-cuda12-pjrt" +version = "0.9.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1b/b8/b725d6ae617b5b6d988eb3b8672e997fff52f1749b3b138c371ae6cf7c2d/jax_cuda12_pjrt-0.9.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:31284cd1c279612cebcf2e9383db9cc2e0703ad8ef56c887476b40331fed7c93", size = 149993391, upload-time = "2026-03-02T11:22:08.125Z" }, + { url = "https://files.pythonhosted.org/packages/a7/2c/8ddb471091b46de99bba7eaa7f4e3983f9c8e74e310e585ff08915ce8b7a/jax_cuda12_pjrt-0.9.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:6447fe3d661b1d824ae7329914c4ac05fe54aeab68c7a7ca23c4f5306dec2e1f", size = 156039796, upload-time = "2026-03-02T11:22:13.316Z" }, +] + +[[package]] +name = "jax-cuda12-plugin" +version = "0.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax-cuda12-pjrt" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/80/0d15d44900bbe2c6276537dc62c8ec49cb08d77b11f8a948b398a3d17eec/jax_cuda12_plugin-0.9.1-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:d1bf3374819c8a87e90df04b7144272e5ac0727a5e1fc4872531c52e387c0e38", size = 5652459, upload-time = "2026-03-02T11:22:22.728Z" }, + { url = "https://files.pythonhosted.org/packages/ca/25/7257a5e4b94634495c138b57f15e77573c440f5128fd44882ae798c70181/jax_cuda12_plugin-0.9.1-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:389c896216a6d4af78709a8f4ca75916ab4ff7bbf4b1af5c09dce9823c003b73", size = 5654488, upload-time = "2026-03-02T11:22:23.978Z" }, + { url = "https://files.pythonhosted.org/packages/45/5b/eadf6199ec61d0640c00c41bb6152c67c2b4b116f7c118ab1eec7034d13e/jax_cuda12_plugin-0.9.1-cp313-cp313t-manylinux_2_27_aarch64.whl", hash = "sha256:6c7a92b7890de3b3cf9f16a017d4ce7f364ef7d3003da86a59d941a2da4a0bda", size = 5665206, upload-time = "2026-03-02T11:22:26.471Z" }, + { url = "https://files.pythonhosted.org/packages/8d/30/302d9331fe44db041efe042d70fa963c88eaa32e92b4c241a2229a6fea76/jax_cuda12_plugin-0.9.1-cp313-cp313t-manylinux_2_27_x86_64.whl", hash = "sha256:0a5d2c6604c8fec04df575190f13908a4aa588fb715c338839099d7cb8f90ad5", size = 5663128, upload-time = "2026-03-02T11:22:27.809Z" }, + { url = "https://files.pythonhosted.org/packages/cc/8a/d626849b8667c971a95784cf468d272d92823133dbb2311ceadcb07427aa/jax_cuda12_plugin-0.9.1-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:d04ac5cff48d18f8497d99c25a26f15bb2df3ba287a4e5eb86c61c89f058c2f9", size = 5653063, upload-time = "2026-03-02T11:22:28.998Z" }, + { url = "https://files.pythonhosted.org/packages/17/e9/3632d7eddb8f282b50bf7c095bba9de91aae0de9baea56d1699d982fe5f8/jax_cuda12_plugin-0.9.1-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:55888255c325f98af961cf5f665ca8223d0e36363eb544fc0286277be2324e8f", size = 5655068, upload-time = "2026-03-02T11:22:30.251Z" }, + { url = "https://files.pythonhosted.org/packages/a9/90/cfa3222b419194524d28b63b3c6d09ce22ddf823701ae73426df6688f26a/jax_cuda12_plugin-0.9.1-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:6b4533951bca6e03b07bb092241f2b471d7675e6c0fd5a4c7cc371d030e0c95a", size = 5665690, upload-time = "2026-03-02T11:22:31.438Z" }, + { url = "https://files.pythonhosted.org/packages/11/65/c123fd6343008df4366840a68f9d5a5c2ee1c7a88b94f0485d1e5bbfd1f8/jax_cuda12_plugin-0.9.1-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:88d3714a638e38f80fa91ba4b9ffc4fd767f4c158ef156067801c74b38d6370f", size = 5663278, upload-time = "2026-03-02T11:22:32.687Z" }, +] + +[package.optional-dependencies] +with-cuda = [ + { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvcc-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "sys_platform == 'linux'" }, +] + +[[package]] +name = "jaxlib" +version = "0.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "scipy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/18/fee700125fe4367c75be1d0f300d13069f5ed119a635ea9199de4b4bc9dc/jaxlib-0.9.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e9915bcaa9ffefd40cd3fdb08a83b16b79f1f3c9ba187884f5b442ad2a47ffd1", size = 57982624, upload-time = "2026-03-02T11:23:31.412Z" }, + { url = "https://files.pythonhosted.org/packages/fd/5f/d4a79d6802f3cef02773852453d9528569dd0896964117d4401658828aba/jaxlib-0.9.1-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:9e88c35248b37d5219423ff8ddca60c6a561e665ded5c4fcbc61f0763e03f1e3", size = 76828438, upload-time = "2026-03-02T11:23:34.793Z" }, + { url = "https://files.pythonhosted.org/packages/3e/2e/d84cafbd07e8cdc7701d9f840f4eea0cfcf3487a99ada14507702172da14/jaxlib-0.9.1-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:da60d967b4ac2084a3e3535ad982392894dd6bdf79c9a56978aba08404a58c82", size = 82473711, upload-time = "2026-03-02T11:23:38.356Z" }, + { url = "https://files.pythonhosted.org/packages/45/e6/4d09ec33a5d096c541025272dc31a36aa9d9a5752b37e05193b23c125810/jaxlib-0.9.1-cp313-cp313-win_amd64.whl", hash = "sha256:7ec6e2f43be6e1ae9321efe9a98affcd8acbe0e1fe59aba1d307ba0462752988", size = 62164682, upload-time = "2026-03-02T11:23:41.761Z" }, + { url = "https://files.pythonhosted.org/packages/8a/be/7d810371aa3bdf30882df60965c15773b8990c90e350a650e366e6dedbaa/jaxlib-0.9.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:872e5917ad20cfde85ce6d50a6dffb205ce551d5c691532f0f07e30c34bbb6c3", size = 58092440, upload-time = "2026-03-02T11:23:46.233Z" }, + { url = "https://files.pythonhosted.org/packages/e9/63/0f5acacd3bd6906f2e1f730ceeafac4afc5cc612f43be4820785608cb951/jaxlib-0.9.1-cp313-cp313t-manylinux_2_27_aarch64.whl", hash = "sha256:469f08a30f6b541557e29c5de61ea6df16ac0ef9225879373bb2b332f1b27d14", size = 76949185, upload-time = "2026-03-02T11:23:49.378Z" }, + { url = "https://files.pythonhosted.org/packages/91/c5/a4dee13627d913c7bd0cf29b7f5c1d6a2605760d08a7cff952f9098ebb61/jaxlib-0.9.1-cp313-cp313t-manylinux_2_27_x86_64.whl", hash = "sha256:2e2225b80689610cbb472822dadf7cc200aa4bdac813112a3f6e074d96b1458c", size = 82584273, upload-time = "2026-03-02T11:23:52.762Z" }, + { url = "https://files.pythonhosted.org/packages/a4/b0/f2c9caa6f545d4ecc1eab528c68c9191e40087f1bc79a6da2e29c6416510/jaxlib-0.9.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:3071bf493f6f48207c56b1e9a5bf895e2acebc5bd40f6f35458e76eb8bf210c7", size = 57984052, upload-time = "2026-03-02T11:23:55.766Z" }, + { url = "https://files.pythonhosted.org/packages/0c/e7/237ec5f4cd07420ef50d79a048b769664dbe306e31bdb10f9dcb9accabe9/jaxlib-0.9.1-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:531dff9fae7aea14449ee544cc1415880cc8a346a9287d347dbd1b2b51d8aabd", size = 76846925, upload-time = "2026-03-02T11:23:59.18Z" }, + { url = "https://files.pythonhosted.org/packages/76/fe/67d2c414b0860d42f4a20b1fadbe7aeffb1b3d885efebd7aedf22a4bc2a2/jaxlib-0.9.1-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:2287a1c891b152c52eb9b73925f57cde01be35d2bab4dad9673d3c83c5982ca8", size = 82484342, upload-time = "2026-03-02T11:24:02.541Z" }, + { url = "https://files.pythonhosted.org/packages/54/0d/a8e27c1c434e489883c1182bd52de27775b8a78013de62e6eabf80991df5/jaxlib-0.9.1-cp314-cp314-win_amd64.whl", hash = "sha256:61160d686e6a4703ef30a6a3aa199c934e6359f42d0aa1c0f9c475d3953b9459", size = 64553355, upload-time = "2026-03-02T11:24:05.976Z" }, + { url = "https://files.pythonhosted.org/packages/fa/4a/e5cb3a32320da2e9496c66045a4e19e16597c92a6496dd493b630585c219/jaxlib-0.9.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5ac3db6b164a8a5b473c77ad9da4f43937d309a27f5cb2f38932930b26e42c68", size = 58096335, upload-time = "2026-03-02T11:24:09.01Z" }, + { url = "https://files.pythonhosted.org/packages/50/d2/35ecc2e92065ac035a954fcb4b752baa72747dcc3a3466525c42c4404958/jaxlib-0.9.1-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:30fe58e8e4e105dffe364a6f0dccca16d93433576d4a015babc83339ca7f1f38", size = 76948543, upload-time = "2026-03-02T11:24:12.026Z" }, + { url = "https://files.pythonhosted.org/packages/ba/cb/a8de776aee88f42937d07472953cf7980e45f5fb30aa9d5ee652b4acc771/jaxlib-0.9.1-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:6b6654a20d54e7cc77d1d54c33f1db851ef9d70bb112b627776178221036e720", size = 82585090, upload-time = "2026-03-02T11:24:15.783Z" }, +] + +[[package]] +name = "jaxtyping" +version = "0.3.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/38/40/a2ea3ce0e3e5f540eb970de7792c90fa58fef1b27d34c83f9fa94fea4729/jaxtyping-0.3.7.tar.gz", hash = "sha256:3bd7d9beb7d3cb01a89f93f90581c6f4fff3e5c5dc3c9307e8f8687a040d10c4", size = 45721, upload-time = "2026-01-30T14:18:47.409Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/42/caf65e9a0576a3abadc537e2f831701ba9081f21317fb3be87d64451587a/jaxtyping-0.3.7-py3-none-any.whl", hash = "sha256:303ab8599edf412eeb40bf06c863e3168fa186cf0e7334703fa741ddd7046e66", size = 56101, upload-time = "2026-01-30T14:18:45.954Z" }, +] + +[[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 = "kiwisolver" +version = "1.4.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5c/3c/85844f1b0feb11ee581ac23fe5fce65cd049a200c1446708cc1b7f922875/kiwisolver-1.4.9.tar.gz", hash = "sha256:c3b22c26c6fd6811b0ae8363b95ca8ce4ea3c202d3d0975b2914310ceb1bcc4d", size = 97564, upload-time = "2025-08-10T21:27:49.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/31/c1/c2686cda909742ab66c7388e9a1a8521a59eb89f8bcfbee28fc980d07e24/kiwisolver-1.4.9-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:a5d0432ccf1c7ab14f9949eec60c5d1f924f17c037e9f8b33352fa05799359b8", size = 123681, upload-time = "2025-08-10T21:26:26.725Z" }, + { url = "https://files.pythonhosted.org/packages/ca/f0/f44f50c9f5b1a1860261092e3bc91ecdc9acda848a8b8c6abfda4a24dd5c/kiwisolver-1.4.9-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:efb3a45b35622bb6c16dbfab491a8f5a391fe0e9d45ef32f4df85658232ca0e2", size = 66464, upload-time = "2025-08-10T21:26:27.733Z" }, + { url = "https://files.pythonhosted.org/packages/2d/7a/9d90a151f558e29c3936b8a47ac770235f436f2120aca41a6d5f3d62ae8d/kiwisolver-1.4.9-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1a12cf6398e8a0a001a059747a1cbf24705e18fe413bc22de7b3d15c67cffe3f", size = 64961, upload-time = "2025-08-10T21:26:28.729Z" }, + { url = "https://files.pythonhosted.org/packages/e9/e9/f218a2cb3a9ffbe324ca29a9e399fa2d2866d7f348ec3a88df87fc248fc5/kiwisolver-1.4.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b67e6efbf68e077dd71d1a6b37e43e1a99d0bff1a3d51867d45ee8908b931098", size = 1474607, upload-time = "2025-08-10T21:26:29.798Z" }, + { url = "https://files.pythonhosted.org/packages/d9/28/aac26d4c882f14de59041636292bc838db8961373825df23b8eeb807e198/kiwisolver-1.4.9-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5656aa670507437af0207645273ccdfee4f14bacd7f7c67a4306d0dcaeaf6eed", size = 1276546, upload-time = "2025-08-10T21:26:31.401Z" }, + { url = "https://files.pythonhosted.org/packages/8b/ad/8bfc1c93d4cc565e5069162f610ba2f48ff39b7de4b5b8d93f69f30c4bed/kiwisolver-1.4.9-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:bfc08add558155345129c7803b3671cf195e6a56e7a12f3dde7c57d9b417f525", size = 1294482, upload-time = "2025-08-10T21:26:32.721Z" }, + { url = "https://files.pythonhosted.org/packages/da/f1/6aca55ff798901d8ce403206d00e033191f63d82dd708a186e0ed2067e9c/kiwisolver-1.4.9-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:40092754720b174e6ccf9e845d0d8c7d8e12c3d71e7fc35f55f3813e96376f78", size = 1343720, upload-time = "2025-08-10T21:26:34.032Z" }, + { url = "https://files.pythonhosted.org/packages/d1/91/eed031876c595c81d90d0f6fc681ece250e14bf6998c3d7c419466b523b7/kiwisolver-1.4.9-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:497d05f29a1300d14e02e6441cf0f5ee81c1ff5a304b0d9fb77423974684e08b", size = 2224907, upload-time = "2025-08-10T21:26:35.824Z" }, + { url = "https://files.pythonhosted.org/packages/e9/ec/4d1925f2e49617b9cca9c34bfa11adefad49d00db038e692a559454dfb2e/kiwisolver-1.4.9-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:bdd1a81a1860476eb41ac4bc1e07b3f07259e6d55bbf739b79c8aaedcf512799", size = 2321334, upload-time = "2025-08-10T21:26:37.534Z" }, + { url = "https://files.pythonhosted.org/packages/43/cb/450cd4499356f68802750c6ddc18647b8ea01ffa28f50d20598e0befe6e9/kiwisolver-1.4.9-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:e6b93f13371d341afee3be9f7c5964e3fe61d5fa30f6a30eb49856935dfe4fc3", size = 2488313, upload-time = "2025-08-10T21:26:39.191Z" }, + { url = "https://files.pythonhosted.org/packages/71/67/fc76242bd99f885651128a5d4fa6083e5524694b7c88b489b1b55fdc491d/kiwisolver-1.4.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:d75aa530ccfaa593da12834b86a0724f58bff12706659baa9227c2ccaa06264c", size = 2291970, upload-time = "2025-08-10T21:26:40.828Z" }, + { url = "https://files.pythonhosted.org/packages/75/bd/f1a5d894000941739f2ae1b65a32892349423ad49c2e6d0771d0bad3fae4/kiwisolver-1.4.9-cp313-cp313-win_amd64.whl", hash = "sha256:dd0a578400839256df88c16abddf9ba14813ec5f21362e1fe65022e00c883d4d", size = 73894, upload-time = "2025-08-10T21:26:42.33Z" }, + { url = "https://files.pythonhosted.org/packages/95/38/dce480814d25b99a391abbddadc78f7c117c6da34be68ca8b02d5848b424/kiwisolver-1.4.9-cp313-cp313-win_arm64.whl", hash = "sha256:d4188e73af84ca82468f09cadc5ac4db578109e52acb4518d8154698d3a87ca2", size = 64995, upload-time = "2025-08-10T21:26:43.889Z" }, + { url = "https://files.pythonhosted.org/packages/e2/37/7d218ce5d92dadc5ebdd9070d903e0c7cf7edfe03f179433ac4d13ce659c/kiwisolver-1.4.9-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:5a0f2724dfd4e3b3ac5a82436a8e6fd16baa7d507117e4279b660fe8ca38a3a1", size = 126510, upload-time = "2025-08-10T21:26:44.915Z" }, + { url = "https://files.pythonhosted.org/packages/23/b0/e85a2b48233daef4b648fb657ebbb6f8367696a2d9548a00b4ee0eb67803/kiwisolver-1.4.9-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:1b11d6a633e4ed84fc0ddafd4ebfd8ea49b3f25082c04ad12b8315c11d504dc1", size = 67903, upload-time = "2025-08-10T21:26:45.934Z" }, + { url = "https://files.pythonhosted.org/packages/44/98/f2425bc0113ad7de24da6bb4dae1343476e95e1d738be7c04d31a5d037fd/kiwisolver-1.4.9-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:61874cdb0a36016354853593cffc38e56fc9ca5aa97d2c05d3dcf6922cd55a11", size = 66402, upload-time = "2025-08-10T21:26:47.101Z" }, + { url = "https://files.pythonhosted.org/packages/98/d8/594657886df9f34c4177cc353cc28ca7e6e5eb562d37ccc233bff43bbe2a/kiwisolver-1.4.9-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:60c439763a969a6af93b4881db0eed8fadf93ee98e18cbc35bc8da868d0c4f0c", size = 1582135, upload-time = "2025-08-10T21:26:48.665Z" }, + { url = "https://files.pythonhosted.org/packages/5c/c6/38a115b7170f8b306fc929e166340c24958347308ea3012c2b44e7e295db/kiwisolver-1.4.9-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:92a2f997387a1b79a75e7803aa7ded2cfbe2823852ccf1ba3bcf613b62ae3197", size = 1389409, upload-time = "2025-08-10T21:26:50.335Z" }, + { url = "https://files.pythonhosted.org/packages/bf/3b/e04883dace81f24a568bcee6eb3001da4ba05114afa622ec9b6fafdc1f5e/kiwisolver-1.4.9-cp313-cp313t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a31d512c812daea6d8b3be3b2bfcbeb091dbb09177706569bcfc6240dcf8b41c", size = 1401763, upload-time = "2025-08-10T21:26:51.867Z" }, + { url = "https://files.pythonhosted.org/packages/9f/80/20ace48e33408947af49d7d15c341eaee69e4e0304aab4b7660e234d6288/kiwisolver-1.4.9-cp313-cp313t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:52a15b0f35dad39862d376df10c5230155243a2c1a436e39eb55623ccbd68185", size = 1453643, upload-time = "2025-08-10T21:26:53.592Z" }, + { url = "https://files.pythonhosted.org/packages/64/31/6ce4380a4cd1f515bdda976a1e90e547ccd47b67a1546d63884463c92ca9/kiwisolver-1.4.9-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:a30fd6fdef1430fd9e1ba7b3398b5ee4e2887783917a687d86ba69985fb08748", size = 2330818, upload-time = "2025-08-10T21:26:55.051Z" }, + { url = "https://files.pythonhosted.org/packages/fa/e9/3f3fcba3bcc7432c795b82646306e822f3fd74df0ee81f0fa067a1f95668/kiwisolver-1.4.9-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:cc9617b46837c6468197b5945e196ee9ca43057bb7d9d1ae688101e4e1dddf64", size = 2419963, upload-time = "2025-08-10T21:26:56.421Z" }, + { url = "https://files.pythonhosted.org/packages/99/43/7320c50e4133575c66e9f7dadead35ab22d7c012a3b09bb35647792b2a6d/kiwisolver-1.4.9-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:0ab74e19f6a2b027ea4f845a78827969af45ce790e6cb3e1ebab71bdf9f215ff", size = 2594639, upload-time = "2025-08-10T21:26:57.882Z" }, + { url = "https://files.pythonhosted.org/packages/65/d6/17ae4a270d4a987ef8a385b906d2bdfc9fce502d6dc0d3aea865b47f548c/kiwisolver-1.4.9-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:dba5ee5d3981160c28d5490f0d1b7ed730c22470ff7f6cc26cfcfaacb9896a07", size = 2391741, upload-time = "2025-08-10T21:26:59.237Z" }, + { url = "https://files.pythonhosted.org/packages/2a/8f/8f6f491d595a9e5912971f3f863d81baddccc8a4d0c3749d6a0dd9ffc9df/kiwisolver-1.4.9-cp313-cp313t-win_arm64.whl", hash = "sha256:0749fd8f4218ad2e851e11cc4dc05c7cbc0cbc4267bdfdb31782e65aace4ee9c", size = 68646, upload-time = "2025-08-10T21:27:00.52Z" }, + { url = "https://files.pythonhosted.org/packages/6b/32/6cc0fbc9c54d06c2969faa9c1d29f5751a2e51809dd55c69055e62d9b426/kiwisolver-1.4.9-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:9928fe1eb816d11ae170885a74d074f57af3a0d65777ca47e9aeb854a1fba386", size = 123806, upload-time = "2025-08-10T21:27:01.537Z" }, + { url = "https://files.pythonhosted.org/packages/b2/dd/2bfb1d4a4823d92e8cbb420fe024b8d2167f72079b3bb941207c42570bdf/kiwisolver-1.4.9-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:d0005b053977e7b43388ddec89fa567f43d4f6d5c2c0affe57de5ebf290dc552", size = 66605, upload-time = "2025-08-10T21:27:03.335Z" }, + { url = "https://files.pythonhosted.org/packages/f7/69/00aafdb4e4509c2ca6064646cba9cd4b37933898f426756adb2cb92ebbed/kiwisolver-1.4.9-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:2635d352d67458b66fd0667c14cb1d4145e9560d503219034a18a87e971ce4f3", size = 64925, upload-time = "2025-08-10T21:27:04.339Z" }, + { url = "https://files.pythonhosted.org/packages/43/dc/51acc6791aa14e5cb6d8a2e28cefb0dc2886d8862795449d021334c0df20/kiwisolver-1.4.9-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:767c23ad1c58c9e827b649a9ab7809fd5fd9db266a9cf02b0e926ddc2c680d58", size = 1472414, upload-time = "2025-08-10T21:27:05.437Z" }, + { url = "https://files.pythonhosted.org/packages/3d/bb/93fa64a81db304ac8a246f834d5094fae4b13baf53c839d6bb6e81177129/kiwisolver-1.4.9-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:72d0eb9fba308b8311685c2268cf7d0a0639a6cd027d8128659f72bdd8a024b4", size = 1281272, upload-time = "2025-08-10T21:27:07.063Z" }, + { url = "https://files.pythonhosted.org/packages/70/e6/6df102916960fb8d05069d4bd92d6d9a8202d5a3e2444494e7cd50f65b7a/kiwisolver-1.4.9-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f68e4f3eeca8fb22cc3d731f9715a13b652795ef657a13df1ad0c7dc0e9731df", size = 1298578, upload-time = "2025-08-10T21:27:08.452Z" }, + { url = "https://files.pythonhosted.org/packages/7c/47/e142aaa612f5343736b087864dbaebc53ea8831453fb47e7521fa8658f30/kiwisolver-1.4.9-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d84cd4061ae292d8ac367b2c3fa3aad11cb8625a95d135fe93f286f914f3f5a6", size = 1345607, upload-time = "2025-08-10T21:27:10.125Z" }, + { url = "https://files.pythonhosted.org/packages/54/89/d641a746194a0f4d1a3670fb900d0dbaa786fb98341056814bc3f058fa52/kiwisolver-1.4.9-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a60ea74330b91bd22a29638940d115df9dc00af5035a9a2a6ad9399ffb4ceca5", size = 2230150, upload-time = "2025-08-10T21:27:11.484Z" }, + { url = "https://files.pythonhosted.org/packages/aa/6b/5ee1207198febdf16ac11f78c5ae40861b809cbe0e6d2a8d5b0b3044b199/kiwisolver-1.4.9-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:ce6a3a4e106cf35c2d9c4fa17c05ce0b180db622736845d4315519397a77beaf", size = 2325979, upload-time = "2025-08-10T21:27:12.917Z" }, + { url = "https://files.pythonhosted.org/packages/fc/ff/b269eefd90f4ae14dcc74973d5a0f6d28d3b9bb1afd8c0340513afe6b39a/kiwisolver-1.4.9-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:77937e5e2a38a7b48eef0585114fe7930346993a88060d0bf886086d2aa49ef5", size = 2491456, upload-time = "2025-08-10T21:27:14.353Z" }, + { url = "https://files.pythonhosted.org/packages/fc/d4/10303190bd4d30de547534601e259a4fbf014eed94aae3e5521129215086/kiwisolver-1.4.9-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:24c175051354f4a28c5d6a31c93906dc653e2bf234e8a4bbfb964892078898ce", size = 2294621, upload-time = "2025-08-10T21:27:15.808Z" }, + { url = "https://files.pythonhosted.org/packages/28/e0/a9a90416fce5c0be25742729c2ea52105d62eda6c4be4d803c2a7be1fa50/kiwisolver-1.4.9-cp314-cp314-win_amd64.whl", hash = "sha256:0763515d4df10edf6d06a3c19734e2566368980d21ebec439f33f9eb936c07b7", size = 75417, upload-time = "2025-08-10T21:27:17.436Z" }, + { url = "https://files.pythonhosted.org/packages/1f/10/6949958215b7a9a264299a7db195564e87900f709db9245e4ebdd3c70779/kiwisolver-1.4.9-cp314-cp314-win_arm64.whl", hash = "sha256:0e4e2bf29574a6a7b7f6cb5fa69293b9f96c928949ac4a53ba3f525dffb87f9c", size = 66582, upload-time = "2025-08-10T21:27:18.436Z" }, + { url = "https://files.pythonhosted.org/packages/ec/79/60e53067903d3bc5469b369fe0dfc6b3482e2133e85dae9daa9527535991/kiwisolver-1.4.9-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:d976bbb382b202f71c67f77b0ac11244021cfa3f7dfd9e562eefcea2df711548", size = 126514, upload-time = "2025-08-10T21:27:19.465Z" }, + { url = "https://files.pythonhosted.org/packages/25/d1/4843d3e8d46b072c12a38c97c57fab4608d36e13fe47d47ee96b4d61ba6f/kiwisolver-1.4.9-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:2489e4e5d7ef9a1c300a5e0196e43d9c739f066ef23270607d45aba368b91f2d", size = 67905, upload-time = "2025-08-10T21:27:20.51Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ae/29ffcbd239aea8b93108de1278271ae764dfc0d803a5693914975f200596/kiwisolver-1.4.9-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:e2ea9f7ab7fbf18fffb1b5434ce7c69a07582f7acc7717720f1d69f3e806f90c", size = 66399, upload-time = "2025-08-10T21:27:21.496Z" }, + { url = "https://files.pythonhosted.org/packages/a1/ae/d7ba902aa604152c2ceba5d352d7b62106bedbccc8e95c3934d94472bfa3/kiwisolver-1.4.9-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b34e51affded8faee0dfdb705416153819d8ea9250bbbf7ea1b249bdeb5f1122", size = 1582197, upload-time = "2025-08-10T21:27:22.604Z" }, + { url = "https://files.pythonhosted.org/packages/f2/41/27c70d427eddb8bc7e4f16420a20fefc6f480312122a59a959fdfe0445ad/kiwisolver-1.4.9-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d8aacd3d4b33b772542b2e01beb50187536967b514b00003bdda7589722d2a64", size = 1390125, upload-time = "2025-08-10T21:27:24.036Z" }, + { url = "https://files.pythonhosted.org/packages/41/42/b3799a12bafc76d962ad69083f8b43b12bf4fe78b097b12e105d75c9b8f1/kiwisolver-1.4.9-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7cf974dd4e35fa315563ac99d6287a1024e4dc2077b8a7d7cd3d2fb65d283134", size = 1402612, upload-time = "2025-08-10T21:27:25.773Z" }, + { url = "https://files.pythonhosted.org/packages/d2/b5/a210ea073ea1cfaca1bb5c55a62307d8252f531beb364e18aa1e0888b5a0/kiwisolver-1.4.9-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:85bd218b5ecfbee8c8a82e121802dcb519a86044c9c3b2e4aef02fa05c6da370", size = 1453990, upload-time = "2025-08-10T21:27:27.089Z" }, + { url = "https://files.pythonhosted.org/packages/5f/ce/a829eb8c033e977d7ea03ed32fb3c1781b4fa0433fbadfff29e39c676f32/kiwisolver-1.4.9-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:0856e241c2d3df4efef7c04a1e46b1936b6120c9bcf36dd216e3acd84bc4fb21", size = 2331601, upload-time = "2025-08-10T21:27:29.343Z" }, + { url = "https://files.pythonhosted.org/packages/e0/4b/b5e97eb142eb9cd0072dacfcdcd31b1c66dc7352b0f7c7255d339c0edf00/kiwisolver-1.4.9-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:9af39d6551f97d31a4deebeac6f45b156f9755ddc59c07b402c148f5dbb6482a", size = 2422041, upload-time = "2025-08-10T21:27:30.754Z" }, + { url = "https://files.pythonhosted.org/packages/40/be/8eb4cd53e1b85ba4edc3a9321666f12b83113a178845593307a3e7891f44/kiwisolver-1.4.9-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:bb4ae2b57fc1d8cbd1cf7b1d9913803681ffa903e7488012be5b76dedf49297f", size = 2594897, upload-time = "2025-08-10T21:27:32.803Z" }, + { url = "https://files.pythonhosted.org/packages/99/dd/841e9a66c4715477ea0abc78da039832fbb09dac5c35c58dc4c41a407b8a/kiwisolver-1.4.9-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:aedff62918805fb62d43a4aa2ecd4482c380dc76cd31bd7c8878588a61bd0369", size = 2391835, upload-time = "2025-08-10T21:27:34.23Z" }, + { url = "https://files.pythonhosted.org/packages/0c/28/4b2e5c47a0da96896fdfdb006340ade064afa1e63675d01ea5ac222b6d52/kiwisolver-1.4.9-cp314-cp314t-win_amd64.whl", hash = "sha256:1fa333e8b2ce4d9660f2cda9c0e1b6bafcfb2457a9d259faa82289e73ec24891", size = 79988, upload-time = "2025-08-10T21:27:35.587Z" }, + { url = "https://files.pythonhosted.org/packages/80/be/3578e8afd18c88cdf9cb4cffde75a96d2be38c5a903f1ed0ceec061bd09e/kiwisolver-1.4.9-cp314-cp314t-win_arm64.whl", hash = "sha256:4a48a2ce79d65d363597ef7b567ce3d14d68783d2b2263d98db3d9477805ba32", size = 70260, upload-time = "2025-08-10T21:27:36.606Z" }, +] + +[[package]] +name = "lczero-training" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "anyio" }, + { name = "flax" }, + { name = "graphviz" }, + { name = "jax", extra = ["cuda12"] }, + { name = "jaxlib" }, + { name = "matplotlib" }, + { name = "meson" }, + { name = "mypy" }, + { name = "numpy" }, + { name = "onnxruntime" }, + { name = "onnxruntime-gpu" }, + { name = "optax" }, + { name = "orbax-checkpoint" }, + { name = "protobuf" }, + { name = "pybind11" }, + { name = "pytest" }, + { name = "python-dotenv" }, + { name = "requests", extra = ["socks"] }, + { name = "ruff" }, + { name = "tensorboardx" }, + { name = "textual" }, +] + +[package.optional-dependencies] +dev = [ + { name = "mypy" }, + { name = "mypy-extensions" }, + { name = "pathspec" }, + { name = "pytest" }, + { name = "typing-extensions" }, +] + +[package.dev-dependencies] +dev = [ + { name = "grpcio-tools" }, + { name = "mypy-protobuf" }, + { name = "textual-dev" }, + { name = "types-protobuf" }, + { name = "types-requests" }, +] + +[package.metadata] +requires-dist = [ + { name = "anyio", specifier = ">=4.10.0" }, + { name = "flax", specifier = ">=0.11.1" }, + { name = "graphviz", specifier = ">=0.21" }, + { name = "jax", extras = ["cuda12"], specifier = "==0.9.1" }, + { name = "jaxlib", specifier = "==0.9.1" }, + { name = "matplotlib", specifier = ">=3.10.6" }, + { name = "meson", specifier = ">=1.10.1" }, + { name = "mypy", specifier = ">=1.17.1" }, + { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.0.0" }, + { name = "mypy-extensions", marker = "extra == 'dev'", specifier = ">=0.4.0" }, + { name = "numpy", specifier = ">=1.24.0" }, + { name = "onnxruntime", specifier = ">=1.23.1" }, + { name = "onnxruntime-gpu", specifier = ">=1.23.0" }, + { name = "optax", specifier = ">=0.2.5" }, + { name = "orbax-checkpoint", specifier = ">=0.11.23" }, + { name = "pathspec", marker = "extra == 'dev'", specifier = ">=0.11.0" }, + { name = "protobuf", specifier = "==6.33.5" }, + { name = "pybind11", specifier = ">=2.10.0" }, + { name = "pytest", specifier = ">=8.4.1" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, + { name = "python-dotenv", specifier = ">=1.1.1" }, + { name = "requests", extras = ["socks"], specifier = ">=2.32.5" }, + { name = "ruff", specifier = ">=0.15.4" }, + { name = "tensorboardx", specifier = ">=2.6.4" }, + { name = "textual", extras = ["dev"], specifier = ">=0.47.0" }, + { name = "typing-extensions", marker = "extra == 'dev'", specifier = ">=4.0.0" }, +] +provides-extras = ["dev"] + +[package.metadata.requires-dev] +dev = [ + { name = "grpcio-tools", specifier = ">=1.76.0" }, + { name = "mypy-protobuf", specifier = ">=3.6.0" }, + { name = "textual-dev", specifier = ">=1.7.0" }, + { name = "types-protobuf", specifier = ">=6.30.2.20250809" }, + { name = "types-requests", specifier = ">=2.32.4.20250913" }, +] + +[[package]] +name = "librt" +version = "0.7.8" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/24/5f3646ff414285e0f7708fa4e946b9bf538345a41d1c375c439467721a5e/librt-0.7.8.tar.gz", hash = "sha256:1a4ede613941d9c3470b0368be851df6bb78ab218635512d0370b27a277a0862", size = 148323, upload-time = "2026-01-14T12:56:16.876Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a1/fe/b1f9de2829cf7fc7649c1dcd202cfd873837c5cc2fc9e526b0e7f716c3d2/librt-0.7.8-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4c3995abbbb60b3c129490fa985dfe6cac11d88fc3c36eeb4fb1449efbbb04fc", size = 57500, upload-time = "2026-01-14T12:55:21.219Z" }, + { url = "https://files.pythonhosted.org/packages/eb/d4/4a60fbe2e53b825f5d9a77325071d61cd8af8506255067bf0c8527530745/librt-0.7.8-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:44e0c2cbc9bebd074cf2cdbe472ca185e824be4e74b1c63a8e934cea674bebf2", size = 59019, upload-time = "2026-01-14T12:55:22.256Z" }, + { url = "https://files.pythonhosted.org/packages/6a/37/61ff80341ba5159afa524445f2d984c30e2821f31f7c73cf166dcafa5564/librt-0.7.8-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:4d2f1e492cae964b3463a03dc77a7fe8742f7855d7258c7643f0ee32b6651dd3", size = 169015, upload-time = "2026-01-14T12:55:23.24Z" }, + { url = "https://files.pythonhosted.org/packages/1c/86/13d4f2d6a93f181ebf2fc953868826653ede494559da8268023fe567fca3/librt-0.7.8-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:451e7ffcef8f785831fdb791bd69211f47e95dc4c6ddff68e589058806f044c6", size = 178161, upload-time = "2026-01-14T12:55:24.826Z" }, + { url = "https://files.pythonhosted.org/packages/88/26/e24ef01305954fc4d771f1f09f3dd682f9eb610e1bec188ffb719374d26e/librt-0.7.8-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3469e1af9f1380e093ae06bedcbdd11e407ac0b303a56bbe9afb1d6824d4982d", size = 193015, upload-time = "2026-01-14T12:55:26.04Z" }, + { url = "https://files.pythonhosted.org/packages/88/a0/92b6bd060e720d7a31ed474d046a69bd55334ec05e9c446d228c4b806ae3/librt-0.7.8-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f11b300027ce19a34f6d24ebb0a25fd0e24a9d53353225a5c1e6cadbf2916b2e", size = 192038, upload-time = "2026-01-14T12:55:27.208Z" }, + { url = "https://files.pythonhosted.org/packages/06/bb/6f4c650253704279c3a214dad188101d1b5ea23be0606628bc6739456624/librt-0.7.8-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:4adc73614f0d3c97874f02f2c7fd2a27854e7e24ad532ea6b965459c5b757eca", size = 186006, upload-time = "2026-01-14T12:55:28.594Z" }, + { url = "https://files.pythonhosted.org/packages/dc/00/1c409618248d43240cadf45f3efb866837fa77e9a12a71481912135eb481/librt-0.7.8-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:60c299e555f87e4c01b2eca085dfccda1dde87f5a604bb45c2906b8305819a93", size = 206888, upload-time = "2026-01-14T12:55:30.214Z" }, + { url = "https://files.pythonhosted.org/packages/d9/83/b2cfe8e76ff5c1c77f8a53da3d5de62d04b5ebf7cf913e37f8bca43b5d07/librt-0.7.8-cp313-cp313-win32.whl", hash = "sha256:b09c52ed43a461994716082ee7d87618096851319bf695d57ec123f2ab708951", size = 44126, upload-time = "2026-01-14T12:55:31.44Z" }, + { url = "https://files.pythonhosted.org/packages/a9/0b/c59d45de56a51bd2d3a401fc63449c0ac163e4ef7f523ea8b0c0dee86ec5/librt-0.7.8-cp313-cp313-win_amd64.whl", hash = "sha256:f8f4a901a3fa28969d6e4519deceab56c55a09d691ea7b12ca830e2fa3461e34", size = 50262, upload-time = "2026-01-14T12:55:33.01Z" }, + { url = "https://files.pythonhosted.org/packages/fc/b9/973455cec0a1ec592395250c474164c4a58ebf3e0651ee920fef1a2623f1/librt-0.7.8-cp313-cp313-win_arm64.whl", hash = "sha256:43d4e71b50763fcdcf64725ac680d8cfa1706c928b844794a7aa0fa9ac8e5f09", size = 43600, upload-time = "2026-01-14T12:55:34.054Z" }, + { url = "https://files.pythonhosted.org/packages/1a/73/fa8814c6ce2d49c3827829cadaa1589b0bf4391660bd4510899393a23ebc/librt-0.7.8-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:be927c3c94c74b05128089a955fba86501c3b544d1d300282cc1b4bd370cb418", size = 57049, upload-time = "2026-01-14T12:55:35.056Z" }, + { url = "https://files.pythonhosted.org/packages/53/fe/f6c70956da23ea235fd2e3cc16f4f0b4ebdfd72252b02d1164dd58b4e6c3/librt-0.7.8-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:7b0803e9008c62a7ef79058233db7ff6f37a9933b8f2573c05b07ddafa226611", size = 58689, upload-time = "2026-01-14T12:55:36.078Z" }, + { url = "https://files.pythonhosted.org/packages/1f/4d/7a2481444ac5fba63050d9abe823e6bc16896f575bfc9c1e5068d516cdce/librt-0.7.8-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:79feb4d00b2a4e0e05c9c56df707934f41fcb5fe53fd9efb7549068d0495b758", size = 166808, upload-time = "2026-01-14T12:55:37.595Z" }, + { url = "https://files.pythonhosted.org/packages/ac/3c/10901d9e18639f8953f57c8986796cfbf4c1c514844a41c9197cf87cb707/librt-0.7.8-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b9122094e3f24aa759c38f46bd8863433820654927370250f460ae75488b66ea", size = 175614, upload-time = "2026-01-14T12:55:38.756Z" }, + { url = "https://files.pythonhosted.org/packages/db/01/5cbdde0951a5090a80e5ba44e6357d375048123c572a23eecfb9326993a7/librt-0.7.8-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e03bea66af33c95ce3addf87a9bf1fcad8d33e757bc479957ddbc0e4f7207ac", size = 189955, upload-time = "2026-01-14T12:55:39.939Z" }, + { url = "https://files.pythonhosted.org/packages/6a/b4/e80528d2f4b7eaf1d437fcbd6fc6ba4cbeb3e2a0cb9ed5a79f47c7318706/librt-0.7.8-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:f1ade7f31675db00b514b98f9ab9a7698c7282dad4be7492589109471852d398", size = 189370, upload-time = "2026-01-14T12:55:41.057Z" }, + { url = "https://files.pythonhosted.org/packages/c1/ab/938368f8ce31a9787ecd4becb1e795954782e4312095daf8fd22420227c8/librt-0.7.8-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:a14229ac62adcf1b90a15992f1ab9c69ae8b99ffb23cb64a90878a6e8a2f5b81", size = 183224, upload-time = "2026-01-14T12:55:42.328Z" }, + { url = "https://files.pythonhosted.org/packages/3c/10/559c310e7a6e4014ac44867d359ef8238465fb499e7eb31b6bfe3e3f86f5/librt-0.7.8-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5bcaaf624fd24e6a0cb14beac37677f90793a96864c67c064a91458611446e83", size = 203541, upload-time = "2026-01-14T12:55:43.501Z" }, + { url = "https://files.pythonhosted.org/packages/f8/db/a0db7acdb6290c215f343835c6efda5b491bb05c3ddc675af558f50fdba3/librt-0.7.8-cp314-cp314-win32.whl", hash = "sha256:7aa7d5457b6c542ecaed79cec4ad98534373c9757383973e638ccced0f11f46d", size = 40657, upload-time = "2026-01-14T12:55:44.668Z" }, + { url = "https://files.pythonhosted.org/packages/72/e0/4f9bdc2a98a798511e81edcd6b54fe82767a715e05d1921115ac70717f6f/librt-0.7.8-cp314-cp314-win_amd64.whl", hash = "sha256:3d1322800771bee4a91f3b4bd4e49abc7d35e65166821086e5afd1e6c0d9be44", size = 46835, upload-time = "2026-01-14T12:55:45.655Z" }, + { url = "https://files.pythonhosted.org/packages/f9/3d/59c6402e3dec2719655a41ad027a7371f8e2334aa794ed11533ad5f34969/librt-0.7.8-cp314-cp314-win_arm64.whl", hash = "sha256:5363427bc6a8c3b1719f8f3845ea53553d301382928a86e8fab7984426949bce", size = 39885, upload-time = "2026-01-14T12:55:47.138Z" }, + { url = "https://files.pythonhosted.org/packages/4e/9c/2481d80950b83085fb14ba3c595db56330d21bbc7d88a19f20165f3538db/librt-0.7.8-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:ca916919793a77e4a98d4a1701e345d337ce53be4a16620f063191f7322ac80f", size = 59161, upload-time = "2026-01-14T12:55:48.45Z" }, + { url = "https://files.pythonhosted.org/packages/96/79/108df2cfc4e672336765d54e3ff887294c1cc36ea4335c73588875775527/librt-0.7.8-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:54feb7b4f2f6706bb82325e836a01be805770443e2400f706e824e91f6441dde", size = 61008, upload-time = "2026-01-14T12:55:49.527Z" }, + { url = "https://files.pythonhosted.org/packages/46/f2/30179898f9994a5637459d6e169b6abdc982012c0a4b2d4c26f50c06f911/librt-0.7.8-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:39a4c76fee41007070f872b648cc2f711f9abf9a13d0c7162478043377b52c8e", size = 187199, upload-time = "2026-01-14T12:55:50.587Z" }, + { url = "https://files.pythonhosted.org/packages/b4/da/f7563db55cebdc884f518ba3791ad033becc25ff68eb70902b1747dc0d70/librt-0.7.8-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ac9c8a458245c7de80bc1b9765b177055efff5803f08e548dd4bb9ab9a8d789b", size = 198317, upload-time = "2026-01-14T12:55:51.991Z" }, + { url = "https://files.pythonhosted.org/packages/b3/6c/4289acf076ad371471fa86718c30ae353e690d3de6167f7db36f429272f1/librt-0.7.8-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:95b67aa7eff150f075fda09d11f6bfb26edffd300f6ab1666759547581e8f666", size = 210334, upload-time = "2026-01-14T12:55:53.682Z" }, + { url = "https://files.pythonhosted.org/packages/4a/7f/377521ac25b78ac0a5ff44127a0360ee6d5ddd3ce7327949876a30533daa/librt-0.7.8-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:535929b6eff670c593c34ff435d5440c3096f20fa72d63444608a5aef64dd581", size = 211031, upload-time = "2026-01-14T12:55:54.827Z" }, + { url = "https://files.pythonhosted.org/packages/c5/b1/e1e96c3e20b23d00cf90f4aad48f0deb4cdfec2f0ed8380d0d85acf98bbf/librt-0.7.8-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:63937bd0f4d1cb56653dc7ae900d6c52c41f0015e25aaf9902481ee79943b33a", size = 204581, upload-time = "2026-01-14T12:55:56.811Z" }, + { url = "https://files.pythonhosted.org/packages/43/71/0f5d010e92ed9747e14bef35e91b6580533510f1e36a8a09eb79ee70b2f0/librt-0.7.8-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:cf243da9e42d914036fd362ac3fa77d80a41cadcd11ad789b1b5eec4daaf67ca", size = 224731, upload-time = "2026-01-14T12:55:58.175Z" }, + { url = "https://files.pythonhosted.org/packages/22/f0/07fb6ab5c39a4ca9af3e37554f9d42f25c464829254d72e4ebbd81da351c/librt-0.7.8-cp314-cp314t-win32.whl", hash = "sha256:171ca3a0a06c643bd0a2f62a8944e1902c94aa8e5da4db1ea9a8daf872685365", size = 41173, upload-time = "2026-01-14T12:55:59.315Z" }, + { url = "https://files.pythonhosted.org/packages/24/d4/7e4be20993dc6a782639625bd2f97f3c66125c7aa80c82426956811cfccf/librt-0.7.8-cp314-cp314t-win_amd64.whl", hash = "sha256:445b7304145e24c60288a2f172b5ce2ca35c0f81605f5299f3fa567e189d2e32", size = 47668, upload-time = "2026-01-14T12:56:00.261Z" }, + { url = "https://files.pythonhosted.org/packages/fc/85/69f92b2a7b3c0f88ffe107c86b952b397004b5b8ea5a81da3d9c04c04422/librt-0.7.8-cp314-cp314t-win_arm64.whl", hash = "sha256:8766ece9de08527deabcd7cb1b4f1a967a385d26e33e536d6d8913db6ef74f06", size = 40550, upload-time = "2026-01-14T12:56:01.542Z" }, +] + +[[package]] +name = "linkify-it-py" +version = "2.0.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "uc-micro-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2a/ae/bb56c6828e4797ba5a4821eec7c43b8bf40f69cda4d4f5f8c8a2810ec96a/linkify-it-py-2.0.3.tar.gz", hash = "sha256:68cda27e162e9215c17d786649d1da0021a451bdc436ef9e0fa0ba5234b9b048", size = 27946, upload-time = "2024-02-04T14:48:04.179Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/1e/b832de447dee8b582cac175871d2f6c3d5077cc56d5575cadba1fd1cccfa/linkify_it_py-2.0.3-py3-none-any.whl", hash = "sha256:6bcbc417b0ac14323382aef5c5192c0075bf8a9d6b41820a2b66371eac6b6d79", size = 19820, upload-time = "2024-02-04T14:48:02.496Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, +] + +[package.optional-dependencies] +linkify = [ + { name = "linkify-it-py" }, +] + +[[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/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" }, + { url = "https://files.pythonhosted.org/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" }, + { url = "https://files.pythonhosted.org/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" }, + { url = "https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" }, + { url = "https://files.pythonhosted.org/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" }, + { url = "https://files.pythonhosted.org/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" }, + { url = "https://files.pythonhosted.org/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" }, + { url = "https://files.pythonhosted.org/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" }, + { url = "https://files.pythonhosted.org/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" }, + { url = "https://files.pythonhosted.org/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" }, + { url = "https://files.pythonhosted.org/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" }, + { url = "https://files.pythonhosted.org/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" }, + { url = "https://files.pythonhosted.org/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" }, + { url = "https://files.pythonhosted.org/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" }, + { url = "https://files.pythonhosted.org/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" }, + { url = "https://files.pythonhosted.org/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" }, + { url = "https://files.pythonhosted.org/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" }, + { url = "https://files.pythonhosted.org/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" }, + { url = "https://files.pythonhosted.org/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" }, + { url = "https://files.pythonhosted.org/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe", size = 11619, upload-time = "2025-09-27T18:37:06.342Z" }, + { url = "https://files.pythonhosted.org/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026", size = 12029, upload-time = "2025-09-27T18:37:07.213Z" }, + { url = "https://files.pythonhosted.org/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737", size = 24408, upload-time = "2025-09-27T18:37:09.572Z" }, + { url = "https://files.pythonhosted.org/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97", size = 23005, upload-time = "2025-09-27T18:37:10.58Z" }, + { url = "https://files.pythonhosted.org/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d", size = 22048, upload-time = "2025-09-27T18:37:11.547Z" }, + { url = "https://files.pythonhosted.org/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda", size = 23821, upload-time = "2025-09-27T18:37:12.48Z" }, + { url = "https://files.pythonhosted.org/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf", size = 21606, upload-time = "2025-09-27T18:37:13.485Z" }, + { url = "https://files.pythonhosted.org/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe", size = 23043, upload-time = "2025-09-27T18:37:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9", size = 14747, upload-time = "2025-09-27T18:37:15.36Z" }, + { url = "https://files.pythonhosted.org/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581", size = 15341, upload-time = "2025-09-27T18:37:16.496Z" }, + { url = "https://files.pythonhosted.org/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4", size = 14073, upload-time = "2025-09-27T18:37:17.476Z" }, + { url = "https://files.pythonhosted.org/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab", size = 11661, upload-time = "2025-09-27T18:37:18.453Z" }, + { url = "https://files.pythonhosted.org/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175", size = 12069, upload-time = "2025-09-27T18:37:19.332Z" }, + { url = "https://files.pythonhosted.org/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634", size = 25670, upload-time = "2025-09-27T18:37:20.245Z" }, + { url = "https://files.pythonhosted.org/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50", size = 23598, upload-time = "2025-09-27T18:37:21.177Z" }, + { url = "https://files.pythonhosted.org/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e", size = 23261, upload-time = "2025-09-27T18:37:22.167Z" }, + { url = "https://files.pythonhosted.org/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5", size = 24835, upload-time = "2025-09-27T18:37:23.296Z" }, + { url = "https://files.pythonhosted.org/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523", size = 22733, upload-time = "2025-09-27T18:37:24.237Z" }, + { url = "https://files.pythonhosted.org/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc", size = 23672, upload-time = "2025-09-27T18:37:25.271Z" }, + { url = "https://files.pythonhosted.org/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d", size = 14819, upload-time = "2025-09-27T18:37:26.285Z" }, + { url = "https://files.pythonhosted.org/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9", size = 15426, upload-time = "2025-09-27T18:37:27.316Z" }, + { url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, +] + +[[package]] +name = "marshmallow" +version = "3.26.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/79/de6c16cc902f4fc372236926b0ce2ab7845268dcc30fb2fbb7f71b418631/marshmallow-3.26.2.tar.gz", hash = "sha256:bbe2adb5a03e6e3571b573f42527c6fe926e17467833660bebd11593ab8dfd57", size = 222095, upload-time = "2025-12-22T06:53:53.309Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/be/2f/5108cb3ee4ba6501748c4908b908e55f42a5b66245b4cfe0c99326e1ef6e/marshmallow-3.26.2-py3-none-any.whl", hash = "sha256:013fa8a3c4c276c24d26d84ce934dc964e2aa794345a0f8c7e5a7191482c8a73", size = 50964, upload-time = "2025-12-22T06:53:51.801Z" }, +] + +[[package]] +name = "matplotlib" +version = "3.10.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "contourpy" }, + { name = "cycler" }, + { name = "fonttools" }, + { name = "kiwisolver" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "pyparsing" }, + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8a/76/d3c6e3a13fe484ebe7718d14e269c9569c4eb0020a968a327acb3b9a8fe6/matplotlib-3.10.8.tar.gz", hash = "sha256:2299372c19d56bcd35cf05a2738308758d32b9eaed2371898d8f5bd33f084aa3", size = 34806269, upload-time = "2025-12-10T22:56:51.155Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3d/b9/15fd5541ef4f5b9a17eefd379356cf12175fe577424e7b1d80676516031a/matplotlib-3.10.8-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:3f2e409836d7f5ac2f1c013110a4d50b9f7edc26328c108915f9075d7d7a91b6", size = 8261076, upload-time = "2025-12-10T22:55:44.648Z" }, + { url = "https://files.pythonhosted.org/packages/8d/a0/2ba3473c1b66b9c74dc7107c67e9008cb1782edbe896d4c899d39ae9cf78/matplotlib-3.10.8-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:56271f3dac49a88d7fca5060f004d9d22b865f743a12a23b1e937a0be4818ee1", size = 8148794, upload-time = "2025-12-10T22:55:46.252Z" }, + { url = "https://files.pythonhosted.org/packages/75/97/a471f1c3eb1fd6f6c24a31a5858f443891d5127e63a7788678d14e249aea/matplotlib-3.10.8-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a0a7f52498f72f13d4a25ea70f35f4cb60642b466cbb0a9be951b5bc3f45a486", size = 8718474, upload-time = "2025-12-10T22:55:47.864Z" }, + { url = "https://files.pythonhosted.org/packages/01/be/cd478f4b66f48256f42927d0acbcd63a26a893136456cd079c0cc24fbabf/matplotlib-3.10.8-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:646d95230efb9ca614a7a594d4fcacde0ac61d25e37dd51710b36477594963ce", size = 9549637, upload-time = "2025-12-10T22:55:50.048Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7c/8dc289776eae5109e268c4fb92baf870678dc048a25d4ac903683b86d5bf/matplotlib-3.10.8-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f89c151aab2e2e23cb3fe0acad1e8b82841fd265379c4cecd0f3fcb34c15e0f6", size = 9613678, upload-time = "2025-12-10T22:55:52.21Z" }, + { url = "https://files.pythonhosted.org/packages/64/40/37612487cc8a437d4dd261b32ca21fe2d79510fe74af74e1f42becb1bdb8/matplotlib-3.10.8-cp313-cp313-win_amd64.whl", hash = "sha256:e8ea3e2d4066083e264e75c829078f9e149fa119d27e19acd503de65e0b13149", size = 8142686, upload-time = "2025-12-10T22:55:54.253Z" }, + { url = "https://files.pythonhosted.org/packages/66/52/8d8a8730e968185514680c2a6625943f70269509c3dcfc0dcf7d75928cb8/matplotlib-3.10.8-cp313-cp313-win_arm64.whl", hash = "sha256:c108a1d6fa78a50646029cb6d49808ff0fc1330fda87fa6f6250c6b5369b6645", size = 8012917, upload-time = "2025-12-10T22:55:56.268Z" }, + { url = "https://files.pythonhosted.org/packages/b5/27/51fe26e1062f298af5ef66343d8ef460e090a27fea73036c76c35821df04/matplotlib-3.10.8-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:ad3d9833a64cf48cc4300f2b406c3d0f4f4724a91c0bd5640678a6ba7c102077", size = 8305679, upload-time = "2025-12-10T22:55:57.856Z" }, + { url = "https://files.pythonhosted.org/packages/2c/1e/4de865bc591ac8e3062e835f42dd7fe7a93168d519557837f0e37513f629/matplotlib-3.10.8-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:eb3823f11823deade26ce3b9f40dcb4a213da7a670013929f31d5f5ed1055b22", size = 8198336, upload-time = "2025-12-10T22:55:59.371Z" }, + { url = "https://files.pythonhosted.org/packages/c6/cb/2f7b6e75fb4dce87ef91f60cac4f6e34f4c145ab036a22318ec837971300/matplotlib-3.10.8-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d9050fee89a89ed57b4fb2c1bfac9a3d0c57a0d55aed95949eedbc42070fea39", size = 8731653, upload-time = "2025-12-10T22:56:01.032Z" }, + { url = "https://files.pythonhosted.org/packages/46/b3/bd9c57d6ba670a37ab31fb87ec3e8691b947134b201f881665b28cc039ff/matplotlib-3.10.8-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b44d07310e404ba95f8c25aa5536f154c0a8ec473303535949e52eb71d0a1565", size = 9561356, upload-time = "2025-12-10T22:56:02.95Z" }, + { url = "https://files.pythonhosted.org/packages/c0/3d/8b94a481456dfc9dfe6e39e93b5ab376e50998cddfd23f4ae3b431708f16/matplotlib-3.10.8-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:0a33deb84c15ede243aead39f77e990469fff93ad1521163305095b77b72ce4a", size = 9614000, upload-time = "2025-12-10T22:56:05.411Z" }, + { url = "https://files.pythonhosted.org/packages/bd/cd/bc06149fe5585ba800b189a6a654a75f1f127e8aab02fd2be10df7fa500c/matplotlib-3.10.8-cp313-cp313t-win_amd64.whl", hash = "sha256:3a48a78d2786784cc2413e57397981fb45c79e968d99656706018d6e62e57958", size = 8220043, upload-time = "2025-12-10T22:56:07.551Z" }, + { url = "https://files.pythonhosted.org/packages/e3/de/b22cf255abec916562cc04eef457c13e58a1990048de0c0c3604d082355e/matplotlib-3.10.8-cp313-cp313t-win_arm64.whl", hash = "sha256:15d30132718972c2c074cd14638c7f4592bd98719e2308bccea40e0538bc0cb5", size = 8062075, upload-time = "2025-12-10T22:56:09.178Z" }, + { url = "https://files.pythonhosted.org/packages/3c/43/9c0ff7a2f11615e516c3b058e1e6e8f9614ddeca53faca06da267c48345d/matplotlib-3.10.8-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:b53285e65d4fa4c86399979e956235deb900be5baa7fc1218ea67fbfaeaadd6f", size = 8262481, upload-time = "2025-12-10T22:56:10.885Z" }, + { url = "https://files.pythonhosted.org/packages/6f/ca/e8ae28649fcdf039fda5ef554b40a95f50592a3c47e6f7270c9561c12b07/matplotlib-3.10.8-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:32f8dce744be5569bebe789e46727946041199030db8aeb2954d26013a0eb26b", size = 8151473, upload-time = "2025-12-10T22:56:12.377Z" }, + { url = "https://files.pythonhosted.org/packages/f1/6f/009d129ae70b75e88cbe7e503a12a4c0670e08ed748a902c2568909e9eb5/matplotlib-3.10.8-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4cf267add95b1c88300d96ca837833d4112756045364f5c734a2276038dae27d", size = 9553896, upload-time = "2025-12-10T22:56:14.432Z" }, + { url = "https://files.pythonhosted.org/packages/f5/26/4221a741eb97967bc1fd5e4c52b9aa5a91b2f4ec05b59f6def4d820f9df9/matplotlib-3.10.8-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2cf5bd12cecf46908f286d7838b2abc6c91cda506c0445b8223a7c19a00df008", size = 9824193, upload-time = "2025-12-10T22:56:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/1f/f3/3abf75f38605772cf48a9daf5821cd4f563472f38b4b828c6fba6fa6d06e/matplotlib-3.10.8-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:41703cc95688f2516b480f7f339d8851a6035f18e100ee6a32bc0b8536a12a9c", size = 9615444, upload-time = "2025-12-10T22:56:18.155Z" }, + { url = "https://files.pythonhosted.org/packages/93/a5/de89ac80f10b8dc615807ee1133cd99ac74082581196d4d9590bea10690d/matplotlib-3.10.8-cp314-cp314-win_amd64.whl", hash = "sha256:83d282364ea9f3e52363da262ce32a09dfe241e4080dcedda3c0db059d3c1f11", size = 8272719, upload-time = "2025-12-10T22:56:20.366Z" }, + { url = "https://files.pythonhosted.org/packages/69/ce/b006495c19ccc0a137b48083168a37bd056392dee02f87dba0472f2797fe/matplotlib-3.10.8-cp314-cp314-win_arm64.whl", hash = "sha256:2c1998e92cd5999e295a731bcb2911c75f597d937341f3030cc24ef2733d78a8", size = 8144205, upload-time = "2025-12-10T22:56:22.239Z" }, + { url = "https://files.pythonhosted.org/packages/68/d9/b31116a3a855bd313c6fcdb7226926d59b041f26061c6c5b1be66a08c826/matplotlib-3.10.8-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:b5a2b97dbdc7d4f353ebf343744f1d1f1cca8aa8bfddb4262fcf4306c3761d50", size = 8305785, upload-time = "2025-12-10T22:56:24.218Z" }, + { url = "https://files.pythonhosted.org/packages/1e/90/6effe8103f0272685767ba5f094f453784057072f49b393e3ea178fe70a5/matplotlib-3.10.8-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:3f5c3e4da343bba819f0234186b9004faba952cc420fbc522dc4e103c1985908", size = 8198361, upload-time = "2025-12-10T22:56:26.787Z" }, + { url = "https://files.pythonhosted.org/packages/d7/65/a73188711bea603615fc0baecca1061429ac16940e2385433cc778a9d8e7/matplotlib-3.10.8-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f62550b9a30afde8c1c3ae450e5eb547d579dd69b25c2fc7a1c67f934c1717a", size = 9561357, upload-time = "2025-12-10T22:56:28.953Z" }, + { url = "https://files.pythonhosted.org/packages/f4/3d/b5c5d5d5be8ce63292567f0e2c43dde9953d3ed86ac2de0a72e93c8f07a1/matplotlib-3.10.8-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:495672de149445ec1b772ff2c9ede9b769e3cb4f0d0aa7fa730d7f59e2d4e1c1", size = 9823610, upload-time = "2025-12-10T22:56:31.455Z" }, + { url = "https://files.pythonhosted.org/packages/4d/4b/e7beb6bbd49f6bae727a12b270a2654d13c397576d25bd6786e47033300f/matplotlib-3.10.8-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:595ba4d8fe983b88f0eec8c26a241e16d6376fe1979086232f481f8f3f67494c", size = 9614011, upload-time = "2025-12-10T22:56:33.85Z" }, + { url = "https://files.pythonhosted.org/packages/7c/e6/76f2813d31f032e65f6f797e3f2f6e4aab95b65015924b1c51370395c28a/matplotlib-3.10.8-cp314-cp314t-win_amd64.whl", hash = "sha256:25d380fe8b1dc32cf8f0b1b448470a77afb195438bafdf1d858bfb876f3edf7b", size = 8362801, upload-time = "2025-12-10T22:56:36.107Z" }, + { url = "https://files.pythonhosted.org/packages/5d/49/d651878698a0b67f23aa28e17f45a6d6dd3d3f933fa29087fa4ce5947b5a/matplotlib-3.10.8-cp314-cp314t-win_arm64.whl", hash = "sha256:113bb52413ea508ce954a02c10ffd0d565f9c3bc7f2eddc27dfe1731e71c7b5f", size = 8192560, upload-time = "2025-12-10T22:56:38.008Z" }, +] + +[[package]] +name = "mdit-py-plugins" +version = "0.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b2/fd/a756d36c0bfba5f6e39a1cdbdbfdd448dc02692467d83816dff4592a1ebc/mdit_py_plugins-0.5.0.tar.gz", hash = "sha256:f4918cb50119f50446560513a8e311d574ff6aaed72606ddae6d35716fe809c6", size = 44655, upload-time = "2025-08-11T07:25:49.083Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/86/dd6e5db36df29e76c7a7699123569a4a18c1623ce68d826ed96c62643cae/mdit_py_plugins-0.5.0-py3-none-any.whl", hash = "sha256:07a08422fc1936a5d26d146759e9155ea466e842f5ab2f7d2266dd084c8dab1f", size = 57205, upload-time = "2025-08-11T07:25:47.597Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "meson" +version = "1.10.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1b/d3/8c43e758cf456273c32652bb8b7a4ec2d74327d8849856b0b714ad671da7/meson-1.10.1.tar.gz", hash = "sha256:c42296f12db316a4515b9375a5df330f2e751ccdd4f608430d41d7d6210e4317", size = 2413969, upload-time = "2026-01-18T14:45:08.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/d5/582789135863eec7c8c1fa31fbde401b3d5d82dbbb4a0973351a1698f738/meson-1.10.1-py3-none-any.whl", hash = "sha256:fe43d1cc2e6de146fbea78f3a062194bcc0e779efc8a0f0d7c35544dfb86731f", size = 1057724, upload-time = "2026-01-18T14:45:02.584Z" }, +] + +[[package]] +name = "ml-dtypes" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48", size = 676888, upload-time = "2025-11-17T22:31:56.907Z" }, + { url = "https://files.pythonhosted.org/packages/d3/b7/dff378afc2b0d5a7d6cd9d3209b60474d9819d1189d347521e1688a60a53/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b", size = 5036993, upload-time = "2025-11-17T22:31:58.497Z" }, + { url = "https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d", size = 5010956, upload-time = "2025-11-17T22:31:59.931Z" }, + { url = "https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328", size = 212224, upload-time = "2025-11-17T22:32:01.349Z" }, + { url = "https://files.pythonhosted.org/packages/8f/75/dfc3775cb36367816e678f69a7843f6f03bd4e2bcd79941e01ea960a068e/ml_dtypes-0.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175", size = 160798, upload-time = "2025-11-17T22:32:02.864Z" }, + { url = "https://files.pythonhosted.org/packages/4f/74/e9ddb35fd1dd43b1106c20ced3f53c2e8e7fc7598c15638e9f80677f81d4/ml_dtypes-0.5.4-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6", size = 702083, upload-time = "2025-11-17T22:32:04.08Z" }, + { url = "https://files.pythonhosted.org/packages/74/f5/667060b0aed1aa63166b22897fdf16dca9eb704e6b4bbf86848d5a181aa7/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d", size = 5354111, upload-time = "2025-11-17T22:32:05.546Z" }, + { url = "https://files.pythonhosted.org/packages/40/49/0f8c498a28c0efa5f5c95a9e374c83ec1385ca41d0e85e7cf40e5d519a21/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298", size = 5366453, upload-time = "2025-11-17T22:32:07.115Z" }, + { url = "https://files.pythonhosted.org/packages/8c/27/12607423d0a9c6bbbcc780ad19f1f6baa2b68b18ce4bddcdc122c4c68dc9/ml_dtypes-0.5.4-cp313-cp313t-win_amd64.whl", hash = "sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6", size = 225612, upload-time = "2025-11-17T22:32:08.615Z" }, + { url = "https://files.pythonhosted.org/packages/e5/80/5a5929e92c72936d5b19872c5fb8fc09327c1da67b3b68c6a13139e77e20/ml_dtypes-0.5.4-cp313-cp313t-win_arm64.whl", hash = "sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1", size = 164145, upload-time = "2025-11-17T22:32:09.782Z" }, + { url = "https://files.pythonhosted.org/packages/72/4e/1339dc6e2557a344f5ba5590872e80346f76f6cb2ac3dd16e4666e88818c/ml_dtypes-0.5.4-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:2b857d3af6ac0d39db1de7c706e69c7f9791627209c3d6dedbfca8c7e5faec22", size = 673781, upload-time = "2025-11-17T22:32:11.364Z" }, + { url = "https://files.pythonhosted.org/packages/04/f9/067b84365c7e83bda15bba2b06c6ca250ce27b20630b1128c435fb7a09aa/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:805cef3a38f4eafae3a5bf9ebdcdb741d0bcfd9e1bd90eb54abd24f928cd2465", size = 5036145, upload-time = "2025-11-17T22:32:12.783Z" }, + { url = "https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f", size = 5010230, upload-time = "2025-11-17T22:32:14.38Z" }, + { url = "https://files.pythonhosted.org/packages/e9/93/2bfed22d2498c468f6bcd0d9f56b033eaa19f33320389314c19ef6766413/ml_dtypes-0.5.4-cp314-cp314-win_amd64.whl", hash = "sha256:8c6a2dcebd6f3903e05d51960a8058d6e131fe69f952a5397e5dbabc841b6d56", size = 221032, upload-time = "2025-11-17T22:32:15.763Z" }, + { url = "https://files.pythonhosted.org/packages/76/a3/9c912fe6ea747bb10fe2f8f54d027eb265db05dfb0c6335e3e063e74e6e8/ml_dtypes-0.5.4-cp314-cp314-win_arm64.whl", hash = "sha256:5a0f68ca8fd8d16583dfa7793973feb86f2fbb56ce3966daf9c9f748f52a2049", size = 163353, upload-time = "2025-11-17T22:32:16.932Z" }, + { url = "https://files.pythonhosted.org/packages/cd/02/48aa7d84cc30ab4ee37624a2fd98c56c02326785750cd212bc0826c2f15b/ml_dtypes-0.5.4-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:bfc534409c5d4b0bf945af29e5d0ab075eae9eecbb549ff8a29280db822f34f9", size = 702085, upload-time = "2025-11-17T22:32:18.175Z" }, + { url = "https://files.pythonhosted.org/packages/5a/e7/85cb99fe80a7a5513253ec7faa88a65306be071163485e9a626fce1b6e84/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2314892cdc3fcf05e373d76d72aaa15fda9fb98625effa73c1d646f331fcecb7", size = 5355358, upload-time = "2025-11-17T22:32:19.7Z" }, + { url = "https://files.pythonhosted.org/packages/79/2b/a826ba18d2179a56e144aef69e57fb2ab7c464ef0b2111940ee8a3a223a2/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0d2ffd05a2575b1519dc928c0b93c06339eb67173ff53acb00724502cda231cf", size = 5366332, upload-time = "2025-11-17T22:32:21.193Z" }, + { url = "https://files.pythonhosted.org/packages/84/44/f4d18446eacb20ea11e82f133ea8f86e2bf2891785b67d9da8d0ab0ef525/ml_dtypes-0.5.4-cp314-cp314t-win_amd64.whl", hash = "sha256:4381fe2f2452a2d7589689693d3162e876b3ddb0a832cde7a414f8e1adf7eab1", size = 236612, upload-time = "2025-11-17T22:32:22.579Z" }, + { url = "https://files.pythonhosted.org/packages/ad/3f/3d42e9a78fe5edf792a83c074b13b9b770092a4fbf3462872f4303135f09/ml_dtypes-0.5.4-cp314-cp314t-win_arm64.whl", hash = "sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d", size = 168825, upload-time = "2025-11-17T22:32:23.766Z" }, +] + +[[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 = "msgpack" +version = "1.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4d/f2/bfb55a6236ed8725a96b0aa3acbd0ec17588e6a2c3b62a93eb513ed8783f/msgpack-1.1.2.tar.gz", hash = "sha256:3b60763c1373dd60f398488069bcdc703cd08a711477b5d480eecc9f9626f47e", size = 173581, upload-time = "2025-10-08T09:15:56.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6b/31/b46518ecc604d7edf3a4f94cb3bf021fc62aa301f0cb849936968164ef23/msgpack-1.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4efd7b5979ccb539c221a4c4e16aac1a533efc97f3b759bb5a5ac9f6d10383bf", size = 81212, upload-time = "2025-10-08T09:15:14.552Z" }, + { url = "https://files.pythonhosted.org/packages/92/dc/c385f38f2c2433333345a82926c6bfa5ecfff3ef787201614317b58dd8be/msgpack-1.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:42eefe2c3e2af97ed470eec850facbe1b5ad1d6eacdbadc42ec98e7dcf68b4b7", size = 84315, upload-time = "2025-10-08T09:15:15.543Z" }, + { url = "https://files.pythonhosted.org/packages/d3/68/93180dce57f684a61a88a45ed13047558ded2be46f03acb8dec6d7c513af/msgpack-1.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1fdf7d83102bf09e7ce3357de96c59b627395352a4024f6e2458501f158bf999", size = 412721, upload-time = "2025-10-08T09:15:16.567Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ba/459f18c16f2b3fc1a1ca871f72f07d70c07bf768ad0a507a698b8052ac58/msgpack-1.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fac4be746328f90caa3cd4bc67e6fe36ca2bf61d5c6eb6d895b6527e3f05071e", size = 424657, upload-time = "2025-10-08T09:15:17.825Z" }, + { url = "https://files.pythonhosted.org/packages/38/f8/4398c46863b093252fe67368b44edc6c13b17f4e6b0e4929dbf0bdb13f23/msgpack-1.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:fffee09044073e69f2bad787071aeec727183e7580443dfeb8556cbf1978d162", size = 402668, upload-time = "2025-10-08T09:15:19.003Z" }, + { url = "https://files.pythonhosted.org/packages/28/ce/698c1eff75626e4124b4d78e21cca0b4cc90043afb80a507626ea354ab52/msgpack-1.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5928604de9b032bc17f5099496417f113c45bc6bc21b5c6920caf34b3c428794", size = 419040, upload-time = "2025-10-08T09:15:20.183Z" }, + { url = "https://files.pythonhosted.org/packages/67/32/f3cd1667028424fa7001d82e10ee35386eea1408b93d399b09fb0aa7875f/msgpack-1.1.2-cp313-cp313-win32.whl", hash = "sha256:a7787d353595c7c7e145e2331abf8b7ff1e6673a6b974ded96e6d4ec09f00c8c", size = 65037, upload-time = "2025-10-08T09:15:21.416Z" }, + { url = "https://files.pythonhosted.org/packages/74/07/1ed8277f8653c40ebc65985180b007879f6a836c525b3885dcc6448ae6cb/msgpack-1.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:a465f0dceb8e13a487e54c07d04ae3ba131c7c5b95e2612596eafde1dccf64a9", size = 72631, upload-time = "2025-10-08T09:15:22.431Z" }, + { url = "https://files.pythonhosted.org/packages/e5/db/0314e4e2db56ebcf450f277904ffd84a7988b9e5da8d0d61ab2d057df2b6/msgpack-1.1.2-cp313-cp313-win_arm64.whl", hash = "sha256:e69b39f8c0aa5ec24b57737ebee40be647035158f14ed4b40e6f150077e21a84", size = 64118, upload-time = "2025-10-08T09:15:23.402Z" }, + { url = "https://files.pythonhosted.org/packages/22/71/201105712d0a2ff07b7873ed3c220292fb2ea5120603c00c4b634bcdafb3/msgpack-1.1.2-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:e23ce8d5f7aa6ea6d2a2b326b4ba46c985dbb204523759984430db7114f8aa00", size = 81127, upload-time = "2025-10-08T09:15:24.408Z" }, + { url = "https://files.pythonhosted.org/packages/1b/9f/38ff9e57a2eade7bf9dfee5eae17f39fc0e998658050279cbb14d97d36d9/msgpack-1.1.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:6c15b7d74c939ebe620dd8e559384be806204d73b4f9356320632d783d1f7939", size = 84981, upload-time = "2025-10-08T09:15:25.812Z" }, + { url = "https://files.pythonhosted.org/packages/8e/a9/3536e385167b88c2cc8f4424c49e28d49a6fc35206d4a8060f136e71f94c/msgpack-1.1.2-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99e2cb7b9031568a2a5c73aa077180f93dd2e95b4f8d3b8e14a73ae94a9e667e", size = 411885, upload-time = "2025-10-08T09:15:27.22Z" }, + { url = "https://files.pythonhosted.org/packages/2f/40/dc34d1a8d5f1e51fc64640b62b191684da52ca469da9cd74e84936ffa4a6/msgpack-1.1.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:180759d89a057eab503cf62eeec0aa61c4ea1200dee709f3a8e9397dbb3b6931", size = 419658, upload-time = "2025-10-08T09:15:28.4Z" }, + { url = "https://files.pythonhosted.org/packages/3b/ef/2b92e286366500a09a67e03496ee8b8ba00562797a52f3c117aa2b29514b/msgpack-1.1.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:04fb995247a6e83830b62f0b07bf36540c213f6eac8e851166d8d86d83cbd014", size = 403290, upload-time = "2025-10-08T09:15:29.764Z" }, + { url = "https://files.pythonhosted.org/packages/78/90/e0ea7990abea5764e4655b8177aa7c63cdfa89945b6e7641055800f6c16b/msgpack-1.1.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:8e22ab046fa7ede9e36eeb4cfad44d46450f37bb05d5ec482b02868f451c95e2", size = 415234, upload-time = "2025-10-08T09:15:31.022Z" }, + { url = "https://files.pythonhosted.org/packages/72/4e/9390aed5db983a2310818cd7d3ec0aecad45e1f7007e0cda79c79507bb0d/msgpack-1.1.2-cp314-cp314-win32.whl", hash = "sha256:80a0ff7d4abf5fecb995fcf235d4064b9a9a8a40a3ab80999e6ac1e30b702717", size = 66391, upload-time = "2025-10-08T09:15:32.265Z" }, + { url = "https://files.pythonhosted.org/packages/6e/f1/abd09c2ae91228c5f3998dbd7f41353def9eac64253de3c8105efa2082f7/msgpack-1.1.2-cp314-cp314-win_amd64.whl", hash = "sha256:9ade919fac6a3e7260b7f64cea89df6bec59104987cbea34d34a2fa15d74310b", size = 73787, upload-time = "2025-10-08T09:15:33.219Z" }, + { url = "https://files.pythonhosted.org/packages/6a/b0/9d9f667ab48b16ad4115c1935d94023b82b3198064cb84a123e97f7466c1/msgpack-1.1.2-cp314-cp314-win_arm64.whl", hash = "sha256:59415c6076b1e30e563eb732e23b994a61c159cec44deaf584e5cc1dd662f2af", size = 66453, upload-time = "2025-10-08T09:15:34.225Z" }, + { url = "https://files.pythonhosted.org/packages/16/67/93f80545eb1792b61a217fa7f06d5e5cb9e0055bed867f43e2b8e012e137/msgpack-1.1.2-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:897c478140877e5307760b0ea66e0932738879e7aa68144d9b78ea4c8302a84a", size = 85264, upload-time = "2025-10-08T09:15:35.61Z" }, + { url = "https://files.pythonhosted.org/packages/87/1c/33c8a24959cf193966ef11a6f6a2995a65eb066bd681fd085afd519a57ce/msgpack-1.1.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a668204fa43e6d02f89dbe79a30b0d67238d9ec4c5bd8a940fc3a004a47b721b", size = 89076, upload-time = "2025-10-08T09:15:36.619Z" }, + { url = "https://files.pythonhosted.org/packages/fc/6b/62e85ff7193663fbea5c0254ef32f0c77134b4059f8da89b958beb7696f3/msgpack-1.1.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5559d03930d3aa0f3aacb4c42c776af1a2ace2611871c84a75afe436695e6245", size = 435242, upload-time = "2025-10-08T09:15:37.647Z" }, + { url = "https://files.pythonhosted.org/packages/c1/47/5c74ecb4cc277cf09f64e913947871682ffa82b3b93c8dad68083112f412/msgpack-1.1.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:70c5a7a9fea7f036b716191c29047374c10721c389c21e9ffafad04df8c52c90", size = 432509, upload-time = "2025-10-08T09:15:38.794Z" }, + { url = "https://files.pythonhosted.org/packages/24/a4/e98ccdb56dc4e98c929a3f150de1799831c0a800583cde9fa022fa90602d/msgpack-1.1.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f2cb069d8b981abc72b41aea1c580ce92d57c673ec61af4c500153a626cb9e20", size = 415957, upload-time = "2025-10-08T09:15:40.238Z" }, + { url = "https://files.pythonhosted.org/packages/da/28/6951f7fb67bc0a4e184a6b38ab71a92d9ba58080b27a77d3e2fb0be5998f/msgpack-1.1.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d62ce1f483f355f61adb5433ebfd8868c5f078d1a52d042b0a998682b4fa8c27", size = 422910, upload-time = "2025-10-08T09:15:41.505Z" }, + { url = "https://files.pythonhosted.org/packages/f0/03/42106dcded51f0a0b5284d3ce30a671e7bd3f7318d122b2ead66ad289fed/msgpack-1.1.2-cp314-cp314t-win32.whl", hash = "sha256:1d1418482b1ee984625d88aa9585db570180c286d942da463533b238b98b812b", size = 75197, upload-time = "2025-10-08T09:15:42.954Z" }, + { url = "https://files.pythonhosted.org/packages/15/86/d0071e94987f8db59d4eeb386ddc64d0bb9b10820a8d82bcd3e53eeb2da6/msgpack-1.1.2-cp314-cp314t-win_amd64.whl", hash = "sha256:5a46bf7e831d09470ad92dff02b8b1ac92175ca36b087f904a0519857c6be3ff", size = 85772, upload-time = "2025-10-08T09:15:43.954Z" }, + { url = "https://files.pythonhosted.org/packages/81/f2/08ace4142eb281c12701fc3b93a10795e4d4dc7f753911d836675050f886/msgpack-1.1.2-cp314-cp314t-win_arm64.whl", hash = "sha256:d99ef64f349d5ec3293688e91486c5fdb925ed03807f64d98d205d2713c60b46", size = 70868, upload-time = "2025-10-08T09:15:44.959Z" }, +] + +[[package]] +name = "multidict" +version = "6.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1a/c2/c2d94cbe6ac1753f3fc980da97b3d930efe1da3af3c9f5125354436c073d/multidict-6.7.1.tar.gz", hash = "sha256:ec6652a1bee61c53a3e5776b6049172c53b6aaba34f18c9ad04f82712bac623d", size = 102010, upload-time = "2026-01-26T02:46:45.979Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/22/929c141d6c0dba87d3e1d38fbdf1ba8baba86b7776469f2bc2d3227a1e67/multidict-6.7.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:2b41f5fed0ed563624f1c17630cb9941cf2309d4df00e494b551b5f3e3d67a23", size = 76174, upload-time = "2026-01-26T02:44:18.509Z" }, + { url = "https://files.pythonhosted.org/packages/c7/75/bc704ae15fee974f8fccd871305e254754167dce5f9e42d88a2def741a1d/multidict-6.7.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:84e61e3af5463c19b67ced91f6c634effb89ef8bfc5ca0267f954451ed4bb6a2", size = 45116, upload-time = "2026-01-26T02:44:19.745Z" }, + { url = "https://files.pythonhosted.org/packages/79/76/55cd7186f498ed080a18440c9013011eb548f77ae1b297206d030eb1180a/multidict-6.7.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:935434b9853c7c112eee7ac891bc4cb86455aa631269ae35442cb316790c1445", size = 43524, upload-time = "2026-01-26T02:44:21.571Z" }, + { url = "https://files.pythonhosted.org/packages/e9/3c/414842ef8d5a1628d68edee29ba0e5bcf235dbfb3ccd3ea303a7fe8c72ff/multidict-6.7.1-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:432feb25a1cb67fe82a9680b4d65fb542e4635cb3166cd9c01560651ad60f177", size = 249368, upload-time = "2026-01-26T02:44:22.803Z" }, + { url = "https://files.pythonhosted.org/packages/f6/32/befed7f74c458b4a525e60519fe8d87eef72bb1e99924fa2b0f9d97a221e/multidict-6.7.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e82d14e3c948952a1a85503817e038cba5905a3352de76b9a465075d072fba23", size = 256952, upload-time = "2026-01-26T02:44:24.306Z" }, + { url = "https://files.pythonhosted.org/packages/03/d6/c878a44ba877f366630c860fdf74bfb203c33778f12b6ac274936853c451/multidict-6.7.1-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:4cfb48c6ea66c83bcaaf7e4dfa7ec1b6bbcf751b7db85a328902796dfde4c060", size = 240317, upload-time = "2026-01-26T02:44:25.772Z" }, + { url = "https://files.pythonhosted.org/packages/68/49/57421b4d7ad2e9e60e25922b08ceb37e077b90444bde6ead629095327a6f/multidict-6.7.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1d540e51b7e8e170174555edecddbd5538105443754539193e3e1061864d444d", size = 267132, upload-time = "2026-01-26T02:44:27.648Z" }, + { url = "https://files.pythonhosted.org/packages/b7/fe/ec0edd52ddbcea2a2e89e174f0206444a61440b40f39704e64dc807a70bd/multidict-6.7.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:273d23f4b40f3dce4d6c8a821c741a86dec62cded82e1175ba3d99be128147ed", size = 268140, upload-time = "2026-01-26T02:44:29.588Z" }, + { url = "https://files.pythonhosted.org/packages/b0/73/6e1b01cbeb458807aa0831742232dbdd1fa92bfa33f52a3f176b4ff3dc11/multidict-6.7.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d624335fd4fa1c08a53f8b4be7676ebde19cd092b3895c421045ca87895b429", size = 254277, upload-time = "2026-01-26T02:44:30.902Z" }, + { url = "https://files.pythonhosted.org/packages/6a/b2/5fb8c124d7561a4974c342bc8c778b471ebbeb3cc17df696f034a7e9afe7/multidict-6.7.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:12fad252f8b267cc75b66e8fc51b3079604e8d43a75428ffe193cd9e2195dfd6", size = 252291, upload-time = "2026-01-26T02:44:32.31Z" }, + { url = "https://files.pythonhosted.org/packages/5a/96/51d4e4e06bcce92577fcd488e22600bd38e4fd59c20cb49434d054903bd2/multidict-6.7.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:03ede2a6ffbe8ef936b92cb4529f27f42be7f56afcdab5ab739cd5f27fb1cbf9", size = 250156, upload-time = "2026-01-26T02:44:33.734Z" }, + { url = "https://files.pythonhosted.org/packages/db/6b/420e173eec5fba721a50e2a9f89eda89d9c98fded1124f8d5c675f7a0c0f/multidict-6.7.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:90efbcf47dbe33dcf643a1e400d67d59abeac5db07dc3f27d6bdeae497a2198c", size = 249742, upload-time = "2026-01-26T02:44:35.222Z" }, + { url = "https://files.pythonhosted.org/packages/44/a3/ec5b5bd98f306bc2aa297b8c6f11a46714a56b1e6ef5ebda50a4f5d7c5fb/multidict-6.7.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:5c4b9bfc148f5a91be9244d6264c53035c8a0dcd2f51f1c3c6e30e30ebaa1c84", size = 262221, upload-time = "2026-01-26T02:44:36.604Z" }, + { url = "https://files.pythonhosted.org/packages/cd/f7/e8c0d0da0cd1e28d10e624604e1a36bcc3353aaebdfdc3a43c72bc683a12/multidict-6.7.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:401c5a650f3add2472d1d288c26deebc540f99e2fb83e9525007a74cd2116f1d", size = 258664, upload-time = "2026-01-26T02:44:38.008Z" }, + { url = "https://files.pythonhosted.org/packages/52/da/151a44e8016dd33feed44f730bd856a66257c1ee7aed4f44b649fb7edeb3/multidict-6.7.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:97891f3b1b3ffbded884e2916cacf3c6fc87b66bb0dde46f7357404750559f33", size = 249490, upload-time = "2026-01-26T02:44:39.386Z" }, + { url = "https://files.pythonhosted.org/packages/87/af/a3b86bf9630b732897f6fc3f4c4714b90aa4361983ccbdcd6c0339b21b0c/multidict-6.7.1-cp313-cp313-win32.whl", hash = "sha256:e1c5988359516095535c4301af38d8a8838534158f649c05dd1050222321bcb3", size = 41695, upload-time = "2026-01-26T02:44:41.318Z" }, + { url = "https://files.pythonhosted.org/packages/b2/35/e994121b0e90e46134673422dd564623f93304614f5d11886b1b3e06f503/multidict-6.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:960c83bf01a95b12b08fd54324a4eb1d5b52c88932b5cba5d6e712bb3ed12eb5", size = 45884, upload-time = "2026-01-26T02:44:42.488Z" }, + { url = "https://files.pythonhosted.org/packages/ca/61/42d3e5dbf661242a69c97ea363f2d7b46c567da8eadef8890022be6e2ab0/multidict-6.7.1-cp313-cp313-win_arm64.whl", hash = "sha256:563fe25c678aaba333d5399408f5ec3c383ca5b663e7f774dd179a520b8144df", size = 43122, upload-time = "2026-01-26T02:44:43.664Z" }, + { url = "https://files.pythonhosted.org/packages/6d/b3/e6b21c6c4f314bb956016b0b3ef2162590a529b84cb831c257519e7fde44/multidict-6.7.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:c76c4bec1538375dad9d452d246ca5368ad6e1c9039dadcf007ae59c70619ea1", size = 83175, upload-time = "2026-01-26T02:44:44.894Z" }, + { url = "https://files.pythonhosted.org/packages/fb/76/23ecd2abfe0957b234f6c960f4ade497f55f2c16aeb684d4ecdbf1c95791/multidict-6.7.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:57b46b24b5d5ebcc978da4ec23a819a9402b4228b8a90d9c656422b4bdd8a963", size = 48460, upload-time = "2026-01-26T02:44:46.106Z" }, + { url = "https://files.pythonhosted.org/packages/c4/57/a0ed92b23f3a042c36bc4227b72b97eca803f5f1801c1ab77c8a212d455e/multidict-6.7.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e954b24433c768ce78ab7929e84ccf3422e46deb45a4dc9f93438f8217fa2d34", size = 46930, upload-time = "2026-01-26T02:44:47.278Z" }, + { url = "https://files.pythonhosted.org/packages/b5/66/02ec7ace29162e447f6382c495dc95826bf931d3818799bbef11e8f7df1a/multidict-6.7.1-cp313-cp313t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:3bd231490fa7217cc832528e1cd8752a96f0125ddd2b5749390f7c3ec8721b65", size = 242582, upload-time = "2026-01-26T02:44:48.604Z" }, + { url = "https://files.pythonhosted.org/packages/58/18/64f5a795e7677670e872673aca234162514696274597b3708b2c0d276cce/multidict-6.7.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:253282d70d67885a15c8a7716f3a73edf2d635793ceda8173b9ecc21f2fb8292", size = 250031, upload-time = "2026-01-26T02:44:50.544Z" }, + { url = "https://files.pythonhosted.org/packages/c8/ed/e192291dbbe51a8290c5686f482084d31bcd9d09af24f63358c3d42fd284/multidict-6.7.1-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0b4c48648d7649c9335cf1927a8b87fa692de3dcb15faa676c6a6f1f1aabda43", size = 228596, upload-time = "2026-01-26T02:44:51.951Z" }, + { url = "https://files.pythonhosted.org/packages/1e/7e/3562a15a60cf747397e7f2180b0a11dc0c38d9175a650e75fa1b4d325e15/multidict-6.7.1-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:98bc624954ec4d2c7cb074b8eefc2b5d0ce7d482e410df446414355d158fe4ca", size = 257492, upload-time = "2026-01-26T02:44:53.902Z" }, + { url = "https://files.pythonhosted.org/packages/24/02/7d0f9eae92b5249bb50ac1595b295f10e263dd0078ebb55115c31e0eaccd/multidict-6.7.1-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:1b99af4d9eec0b49927b4402bcbb58dea89d3e0db8806a4086117019939ad3dd", size = 255899, upload-time = "2026-01-26T02:44:55.316Z" }, + { url = "https://files.pythonhosted.org/packages/00/e3/9b60ed9e23e64c73a5cde95269ef1330678e9c6e34dd4eb6b431b85b5a10/multidict-6.7.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6aac4f16b472d5b7dc6f66a0d49dd57b0e0902090be16594dc9ebfd3d17c47e7", size = 247970, upload-time = "2026-01-26T02:44:56.783Z" }, + { url = "https://files.pythonhosted.org/packages/3e/06/538e58a63ed5cfb0bd4517e346b91da32fde409d839720f664e9a4ae4f9d/multidict-6.7.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:21f830fe223215dffd51f538e78c172ed7c7f60c9b96a2bf05c4848ad49921c3", size = 245060, upload-time = "2026-01-26T02:44:58.195Z" }, + { url = "https://files.pythonhosted.org/packages/b2/2f/d743a3045a97c895d401e9bd29aaa09b94f5cbdf1bd561609e5a6c431c70/multidict-6.7.1-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:f5dd81c45b05518b9aa4da4aa74e1c93d715efa234fd3e8a179df611cc85e5f4", size = 235888, upload-time = "2026-01-26T02:44:59.57Z" }, + { url = "https://files.pythonhosted.org/packages/38/83/5a325cac191ab28b63c52f14f1131f3b0a55ba3b9aa65a6d0bf2a9b921a0/multidict-6.7.1-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:eb304767bca2bb92fb9c5bd33cedc95baee5bb5f6c88e63706533a1c06ad08c8", size = 243554, upload-time = "2026-01-26T02:45:01.054Z" }, + { url = "https://files.pythonhosted.org/packages/20/1f/9d2327086bd15da2725ef6aae624208e2ef828ed99892b17f60c344e57ed/multidict-6.7.1-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:c9035dde0f916702850ef66460bc4239d89d08df4d02023a5926e7446724212c", size = 252341, upload-time = "2026-01-26T02:45:02.484Z" }, + { url = "https://files.pythonhosted.org/packages/e8/2c/2a1aa0280cf579d0f6eed8ee5211c4f1730bd7e06c636ba2ee6aafda302e/multidict-6.7.1-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:af959b9beeb66c822380f222f0e0a1889331597e81f1ded7f374f3ecb0fd6c52", size = 246391, upload-time = "2026-01-26T02:45:03.862Z" }, + { url = "https://files.pythonhosted.org/packages/e5/03/7ca022ffc36c5a3f6e03b179a5ceb829be9da5783e6fe395f347c0794680/multidict-6.7.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:41f2952231456154ee479651491e94118229844dd7226541788be783be2b5108", size = 243422, upload-time = "2026-01-26T02:45:05.296Z" }, + { url = "https://files.pythonhosted.org/packages/dc/1d/b31650eab6c5778aceed46ba735bd97f7c7d2f54b319fa916c0f96e7805b/multidict-6.7.1-cp313-cp313t-win32.whl", hash = "sha256:df9f19c28adcb40b6aae30bbaa1478c389efd50c28d541d76760199fc1037c32", size = 47770, upload-time = "2026-01-26T02:45:06.754Z" }, + { url = "https://files.pythonhosted.org/packages/ac/5b/2d2d1d522e51285bd61b1e20df8f47ae1a9d80839db0b24ea783b3832832/multidict-6.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:d54ecf9f301853f2c5e802da559604b3e95bb7a3b01a9c295c6ee591b9882de8", size = 53109, upload-time = "2026-01-26T02:45:08.044Z" }, + { url = "https://files.pythonhosted.org/packages/3d/a3/cc409ba012c83ca024a308516703cf339bdc4b696195644a7215a5164a24/multidict-6.7.1-cp313-cp313t-win_arm64.whl", hash = "sha256:5a37ca18e360377cfda1d62f5f382ff41f2b8c4ccb329ed974cc2e1643440118", size = 45573, upload-time = "2026-01-26T02:45:09.349Z" }, + { url = "https://files.pythonhosted.org/packages/91/cc/db74228a8be41884a567e88a62fd589a913708fcf180d029898c17a9a371/multidict-6.7.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8f333ec9c5eb1b7105e3b84b53141e66ca05a19a605368c55450b6ba208cb9ee", size = 75190, upload-time = "2026-01-26T02:45:10.651Z" }, + { url = "https://files.pythonhosted.org/packages/d5/22/492f2246bb5b534abd44804292e81eeaf835388901f0c574bac4eeec73c5/multidict-6.7.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:a407f13c188f804c759fc6a9f88286a565c242a76b27626594c133b82883b5c2", size = 44486, upload-time = "2026-01-26T02:45:11.938Z" }, + { url = "https://files.pythonhosted.org/packages/f1/4f/733c48f270565d78b4544f2baddc2fb2a245e5a8640254b12c36ac7ac68e/multidict-6.7.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0e161ddf326db5577c3a4cc2d8648f81456e8a20d40415541587a71620d7a7d1", size = 43219, upload-time = "2026-01-26T02:45:14.346Z" }, + { url = "https://files.pythonhosted.org/packages/24/bb/2c0c2287963f4259c85e8bcbba9182ced8d7fca65c780c38e99e61629d11/multidict-6.7.1-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:1e3a8bb24342a8201d178c3b4984c26ba81a577c80d4d525727427460a50c22d", size = 245132, upload-time = "2026-01-26T02:45:15.712Z" }, + { url = "https://files.pythonhosted.org/packages/a7/f9/44d4b3064c65079d2467888794dea218d1601898ac50222ab8a9a8094460/multidict-6.7.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97231140a50f5d447d3164f994b86a0bed7cd016e2682f8650d6a9158e14fd31", size = 252420, upload-time = "2026-01-26T02:45:17.293Z" }, + { url = "https://files.pythonhosted.org/packages/8b/13/78f7275e73fa17b24c9a51b0bd9d73ba64bb32d0ed51b02a746eb876abe7/multidict-6.7.1-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:6b10359683bd8806a200fd2909e7c8ca3a7b24ec1d8132e483d58e791d881048", size = 233510, upload-time = "2026-01-26T02:45:19.356Z" }, + { url = "https://files.pythonhosted.org/packages/4b/25/8167187f62ae3cbd52da7893f58cb036b47ea3fb67138787c76800158982/multidict-6.7.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:283ddac99f7ac25a4acadbf004cb5ae34480bbeb063520f70ce397b281859362", size = 264094, upload-time = "2026-01-26T02:45:20.834Z" }, + { url = "https://files.pythonhosted.org/packages/a1/e7/69a3a83b7b030cf283fb06ce074a05a02322359783424d7edf0f15fe5022/multidict-6.7.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:538cec1e18c067d0e6103aa9a74f9e832904c957adc260e61cd9d8cf0c3b3d37", size = 260786, upload-time = "2026-01-26T02:45:22.818Z" }, + { url = "https://files.pythonhosted.org/packages/fe/3b/8ec5074bcfc450fe84273713b4b0a0dd47c0249358f5d82eb8104ffe2520/multidict-6.7.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7eee46ccb30ff48a1e35bb818cc90846c6be2b68240e42a78599166722cea709", size = 248483, upload-time = "2026-01-26T02:45:24.368Z" }, + { url = "https://files.pythonhosted.org/packages/48/5a/d5a99e3acbca0e29c5d9cba8f92ceb15dce78bab963b308ae692981e3a5d/multidict-6.7.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:fa263a02f4f2dd2d11a7b1bb4362aa7cb1049f84a9235d31adf63f30143469a0", size = 248403, upload-time = "2026-01-26T02:45:25.982Z" }, + { url = "https://files.pythonhosted.org/packages/35/48/e58cd31f6c7d5102f2a4bf89f96b9cf7e00b6c6f3d04ecc44417c00a5a3c/multidict-6.7.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:2e1425e2f99ec5bd36c15a01b690a1a2456209c5deed58f95469ffb46039ccbb", size = 240315, upload-time = "2026-01-26T02:45:27.487Z" }, + { url = "https://files.pythonhosted.org/packages/94/33/1cd210229559cb90b6786c30676bb0c58249ff42f942765f88793b41fdce/multidict-6.7.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:497394b3239fc6f0e13a78a3e1b61296e72bf1c5f94b4c4eb80b265c37a131cd", size = 245528, upload-time = "2026-01-26T02:45:28.991Z" }, + { url = "https://files.pythonhosted.org/packages/64/f2/6e1107d226278c876c783056b7db43d800bb64c6131cec9c8dfb6903698e/multidict-6.7.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:233b398c29d3f1b9676b4b6f75c518a06fcb2ea0b925119fb2c1bc35c05e1601", size = 258784, upload-time = "2026-01-26T02:45:30.503Z" }, + { url = "https://files.pythonhosted.org/packages/4d/c1/11f664f14d525e4a1b5327a82d4de61a1db604ab34c6603bb3c2cc63ad34/multidict-6.7.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:93b1818e4a6e0930454f0f2af7dfce69307ca03cdcfb3739bf4d91241967b6c1", size = 251980, upload-time = "2026-01-26T02:45:32.603Z" }, + { url = "https://files.pythonhosted.org/packages/e1/9f/75a9ac888121d0c5bbd4ecf4eead45668b1766f6baabfb3b7f66a410e231/multidict-6.7.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f33dc2a3abe9249ea5d8360f969ec7f4142e7ac45ee7014d8f8d5acddf178b7b", size = 243602, upload-time = "2026-01-26T02:45:34.043Z" }, + { url = "https://files.pythonhosted.org/packages/9a/e7/50bf7b004cc8525d80dbbbedfdc7aed3e4c323810890be4413e589074032/multidict-6.7.1-cp314-cp314-win32.whl", hash = "sha256:3ab8b9d8b75aef9df299595d5388b14530839f6422333357af1339443cff777d", size = 40930, upload-time = "2026-01-26T02:45:36.278Z" }, + { url = "https://files.pythonhosted.org/packages/e0/bf/52f25716bbe93745595800f36fb17b73711f14da59ed0bb2eba141bc9f0f/multidict-6.7.1-cp314-cp314-win_amd64.whl", hash = "sha256:5e01429a929600e7dab7b166062d9bb54a5eed752384c7384c968c2afab8f50f", size = 45074, upload-time = "2026-01-26T02:45:37.546Z" }, + { url = "https://files.pythonhosted.org/packages/97/ab/22803b03285fa3a525f48217963da3a65ae40f6a1b6f6cf2768879e208f9/multidict-6.7.1-cp314-cp314-win_arm64.whl", hash = "sha256:4885cb0e817aef5d00a2e8451d4665c1808378dc27c2705f1bf4ef8505c0d2e5", size = 42471, upload-time = "2026-01-26T02:45:38.889Z" }, + { url = "https://files.pythonhosted.org/packages/e0/6d/f9293baa6146ba9507e360ea0292b6422b016907c393e2f63fc40ab7b7b5/multidict-6.7.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:0458c978acd8e6ea53c81eefaddbbee9c6c5e591f41b3f5e8e194780fe026581", size = 82401, upload-time = "2026-01-26T02:45:40.254Z" }, + { url = "https://files.pythonhosted.org/packages/7a/68/53b5494738d83558d87c3c71a486504d8373421c3e0dbb6d0db48ad42ee0/multidict-6.7.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:c0abd12629b0af3cf590982c0b413b1e7395cd4ec026f30986818ab95bfaa94a", size = 48143, upload-time = "2026-01-26T02:45:41.635Z" }, + { url = "https://files.pythonhosted.org/packages/37/e8/5284c53310dcdc99ce5d66563f6e5773531a9b9fe9ec7a615e9bc306b05f/multidict-6.7.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:14525a5f61d7d0c94b368a42cff4c9a4e7ba2d52e2672a7b23d84dc86fb02b0c", size = 46507, upload-time = "2026-01-26T02:45:42.99Z" }, + { url = "https://files.pythonhosted.org/packages/e4/fc/6800d0e5b3875568b4083ecf5f310dcf91d86d52573160834fb4bfcf5e4f/multidict-6.7.1-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:17307b22c217b4cf05033dabefe68255a534d637c6c9b0cc8382718f87be4262", size = 239358, upload-time = "2026-01-26T02:45:44.376Z" }, + { url = "https://files.pythonhosted.org/packages/41/75/4ad0973179361cdf3a113905e6e088173198349131be2b390f9fa4da5fc6/multidict-6.7.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7a7e590ff876a3eaf1c02a4dfe0724b6e69a9e9de6d8f556816f29c496046e59", size = 246884, upload-time = "2026-01-26T02:45:47.167Z" }, + { url = "https://files.pythonhosted.org/packages/c3/9c/095bb28b5da139bd41fb9a5d5caff412584f377914bd8787c2aa98717130/multidict-6.7.1-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:5fa6a95dfee63893d80a34758cd0e0c118a30b8dcb46372bf75106c591b77889", size = 225878, upload-time = "2026-01-26T02:45:48.698Z" }, + { url = "https://files.pythonhosted.org/packages/07/d0/c0a72000243756e8f5a277b6b514fa005f2c73d481b7d9e47cd4568aa2e4/multidict-6.7.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a0543217a6a017692aa6ae5cc39adb75e587af0f3a82288b1492eb73dd6cc2a4", size = 253542, upload-time = "2026-01-26T02:45:50.164Z" }, + { url = "https://files.pythonhosted.org/packages/c0/6b/f69da15289e384ecf2a68837ec8b5ad8c33e973aa18b266f50fe55f24b8c/multidict-6.7.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f99fe611c312b3c1c0ace793f92464d8cd263cc3b26b5721950d977b006b6c4d", size = 252403, upload-time = "2026-01-26T02:45:51.779Z" }, + { url = "https://files.pythonhosted.org/packages/a2/76/b9669547afa5a1a25cd93eaca91c0da1c095b06b6d2d8ec25b713588d3a1/multidict-6.7.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9004d8386d133b7e6135679424c91b0b854d2d164af6ea3f289f8f2761064609", size = 244889, upload-time = "2026-01-26T02:45:53.27Z" }, + { url = "https://files.pythonhosted.org/packages/7e/a9/a50d2669e506dad33cfc45b5d574a205587b7b8a5f426f2fbb2e90882588/multidict-6.7.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e628ef0e6859ffd8273c69412a2465c4be4a9517d07261b33334b5ec6f3c7489", size = 241982, upload-time = "2026-01-26T02:45:54.919Z" }, + { url = "https://files.pythonhosted.org/packages/c5/bb/1609558ad8b456b4827d3c5a5b775c93b87878fd3117ed3db3423dfbce1b/multidict-6.7.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:841189848ba629c3552035a6a7f5bf3b02eb304e9fea7492ca220a8eda6b0e5c", size = 232415, upload-time = "2026-01-26T02:45:56.981Z" }, + { url = "https://files.pythonhosted.org/packages/d8/59/6f61039d2aa9261871e03ab9dc058a550d240f25859b05b67fd70f80d4b3/multidict-6.7.1-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:ce1bbd7d780bb5a0da032e095c951f7014d6b0a205f8318308140f1a6aba159e", size = 240337, upload-time = "2026-01-26T02:45:58.698Z" }, + { url = "https://files.pythonhosted.org/packages/a1/29/fdc6a43c203890dc2ae9249971ecd0c41deaedfe00d25cb6564b2edd99eb/multidict-6.7.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:b26684587228afed0d50cf804cc71062cc9c1cdf55051c4c6345d372947b268c", size = 248788, upload-time = "2026-01-26T02:46:00.862Z" }, + { url = "https://files.pythonhosted.org/packages/a9/14/a153a06101323e4cf086ecee3faadba52ff71633d471f9685c42e3736163/multidict-6.7.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:9f9af11306994335398293f9958071019e3ab95e9a707dc1383a35613f6abcb9", size = 242842, upload-time = "2026-01-26T02:46:02.824Z" }, + { url = "https://files.pythonhosted.org/packages/41/5f/604ae839e64a4a6efc80db94465348d3b328ee955e37acb24badbcd24d83/multidict-6.7.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:b4938326284c4f1224178a560987b6cf8b4d38458b113d9b8c1db1a836e640a2", size = 240237, upload-time = "2026-01-26T02:46:05.898Z" }, + { url = "https://files.pythonhosted.org/packages/5f/60/c3a5187bf66f6fb546ff4ab8fb5a077cbdd832d7b1908d4365c7f74a1917/multidict-6.7.1-cp314-cp314t-win32.whl", hash = "sha256:98655c737850c064a65e006a3df7c997cd3b220be4ec8fe26215760b9697d4d7", size = 48008, upload-time = "2026-01-26T02:46:07.468Z" }, + { url = "https://files.pythonhosted.org/packages/0c/f7/addf1087b860ac60e6f382240f64fb99f8bfb532bb06f7c542b83c29ca61/multidict-6.7.1-cp314-cp314t-win_amd64.whl", hash = "sha256:497bde6223c212ba11d462853cfa4f0ae6ef97465033e7dc9940cdb3ab5b48e5", size = 53542, upload-time = "2026-01-26T02:46:08.809Z" }, + { url = "https://files.pythonhosted.org/packages/4c/81/4629d0aa32302ef7b2ec65c75a728cc5ff4fa410c50096174c1632e70b3e/multidict-6.7.1-cp314-cp314t-win_arm64.whl", hash = "sha256:2bbd113e0d4af5db41d5ebfe9ccaff89de2120578164f86a5d17d5a576d1e5b2", size = 44719, upload-time = "2026-01-26T02:46:11.146Z" }, + { url = "https://files.pythonhosted.org/packages/81/08/7036c080d7117f28a4af526d794aab6a84463126db031b007717c1a6676e/multidict-6.7.1-py3-none-any.whl", hash = "sha256:55d97cc6dae627efa6a6e548885712d4864b81110ac76fa4e534c03819fa4a56", size = 12319, upload-time = "2026-01-26T02:46:44.004Z" }, +] + +[[package]] +name = "mypy" +version = "1.19.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "librt", marker = "platform_python_implementation != 'PyPy'" }, + { name = "mypy-extensions" }, + { name = "pathspec" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f5/db/4efed9504bc01309ab9c2da7e352cc223569f05478012b5d9ece38fd44d2/mypy-1.19.1.tar.gz", hash = "sha256:19d88bb05303fe63f71dd2c6270daca27cb9401c4ca8255fe50d1d920e0eb9ba", size = 3582404, upload-time = "2025-12-15T05:03:48.42Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/9f/a6abae693f7a0c697dbb435aac52e958dc8da44e92e08ba88d2e42326176/mypy-1.19.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e3157c7594ff2ef1634ee058aafc56a82db665c9438fd41b390f3bde1ab12250", size = 13201927, upload-time = "2025-12-15T05:02:29.138Z" }, + { url = "https://files.pythonhosted.org/packages/9a/a4/45c35ccf6e1c65afc23a069f50e2c66f46bd3798cbe0d680c12d12935caa/mypy-1.19.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bdb12f69bcc02700c2b47e070238f42cb87f18c0bc1fc4cdb4fb2bc5fd7a3b8b", size = 12206730, upload-time = "2025-12-15T05:03:01.325Z" }, + { url = "https://files.pythonhosted.org/packages/05/bb/cdcf89678e26b187650512620eec8368fded4cfd99cfcb431e4cdfd19dec/mypy-1.19.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f859fb09d9583a985be9a493d5cfc5515b56b08f7447759a0c5deaf68d80506e", size = 12724581, upload-time = "2025-12-15T05:03:20.087Z" }, + { url = "https://files.pythonhosted.org/packages/d1/32/dd260d52babf67bad8e6770f8e1102021877ce0edea106e72df5626bb0ec/mypy-1.19.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c9a6538e0415310aad77cb94004ca6482330fece18036b5f360b62c45814c4ef", size = 13616252, upload-time = "2025-12-15T05:02:49.036Z" }, + { url = "https://files.pythonhosted.org/packages/71/d0/5e60a9d2e3bd48432ae2b454b7ef2b62a960ab51292b1eda2a95edd78198/mypy-1.19.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:da4869fc5e7f62a88f3fe0b5c919d1d9f7ea3cef92d3689de2823fd27e40aa75", size = 13840848, upload-time = "2025-12-15T05:02:55.95Z" }, + { url = "https://files.pythonhosted.org/packages/98/76/d32051fa65ecf6cc8c6610956473abdc9b4c43301107476ac03559507843/mypy-1.19.1-cp313-cp313-win_amd64.whl", hash = "sha256:016f2246209095e8eda7538944daa1d60e1e8134d98983b9fc1e92c1fc0cb8dd", size = 10135510, upload-time = "2025-12-15T05:02:58.438Z" }, + { url = "https://files.pythonhosted.org/packages/de/eb/b83e75f4c820c4247a58580ef86fcd35165028f191e7e1ba57128c52782d/mypy-1.19.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:06e6170bd5836770e8104c8fdd58e5e725cfeb309f0a6c681a811f557e97eac1", size = 13199744, upload-time = "2025-12-15T05:03:30.823Z" }, + { url = "https://files.pythonhosted.org/packages/94/28/52785ab7bfa165f87fcbb61547a93f98bb20e7f82f90f165a1f69bce7b3d/mypy-1.19.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:804bd67b8054a85447c8954215a906d6eff9cabeabe493fb6334b24f4bfff718", size = 12215815, upload-time = "2025-12-15T05:02:42.323Z" }, + { url = "https://files.pythonhosted.org/packages/0a/c6/bdd60774a0dbfb05122e3e925f2e9e846c009e479dcec4821dad881f5b52/mypy-1.19.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:21761006a7f497cb0d4de3d8ef4ca70532256688b0523eee02baf9eec895e27b", size = 12740047, upload-time = "2025-12-15T05:03:33.168Z" }, + { url = "https://files.pythonhosted.org/packages/32/2a/66ba933fe6c76bd40d1fe916a83f04fed253152f451a877520b3c4a5e41e/mypy-1.19.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:28902ee51f12e0f19e1e16fbe2f8f06b6637f482c459dd393efddd0ec7f82045", size = 13601998, upload-time = "2025-12-15T05:03:13.056Z" }, + { url = "https://files.pythonhosted.org/packages/e3/da/5055c63e377c5c2418760411fd6a63ee2b96cf95397259038756c042574f/mypy-1.19.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:481daf36a4c443332e2ae9c137dfee878fcea781a2e3f895d54bd3002a900957", size = 13807476, upload-time = "2025-12-15T05:03:17.977Z" }, + { url = "https://files.pythonhosted.org/packages/cd/09/4ebd873390a063176f06b0dbf1f7783dd87bd120eae7727fa4ae4179b685/mypy-1.19.1-cp314-cp314-win_amd64.whl", hash = "sha256:8bb5c6f6d043655e055be9b542aa5f3bdd30e4f3589163e85f93f3640060509f", size = 10281872, upload-time = "2025-12-15T05:03:05.549Z" }, + { url = "https://files.pythonhosted.org/packages/8d/f4/4ce9a05ce5ded1de3ec1c1d96cf9f9504a04e54ce0ed55cfa38619a32b8d/mypy-1.19.1-py3-none-any.whl", hash = "sha256:f1235f5ea01b7db5468d53ece6aaddf1ad0b88d9e7462b86ef96fe04995d7247", size = 2471239, upload-time = "2025-12-15T05:03:07.248Z" }, +] + +[[package]] +name = "mypy-extensions" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, +] + +[[package]] +name = "mypy-protobuf" +version = "5.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf" }, + { name = "types-protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d5/48/658827446368bca30a94e545598065587ece9cd09b678d7d2895c37a59d2/mypy_protobuf-5.0.0.tar.gz", hash = "sha256:6fdd1cfdbb4419c713291d800a332d4bba6510dbd1341ed95e0bcc82fcadb6b5", size = 37309, upload-time = "2026-01-13T17:10:12.616Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/b1/ab1e7a49930a8c1d1f7a570bbd4ec7d552ef035acc7aa4b97906e17a34a9/mypy_protobuf-5.0.0-py3-none-any.whl", hash = "sha256:3a7dd753ef3e3b8783a824eb51f07983f62812f9ec066e4fbb1b22d6c5dc36d0", size = 26008, upload-time = "2026-01-13T17:10:11.053Z" }, +] + +[[package]] +name = "nest-asyncio" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/83/f8/51569ac65d696c8ecbee95938f89d4abf00f47d58d48f6fbabfe8f0baefe/nest_asyncio-1.6.0.tar.gz", hash = "sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe", size = 7418, upload-time = "2024-01-21T14:25:19.227Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/c4/c2971a3ba4c6103a3d10c4b0f24f461ddc027f0f09763220cf35ca1401b3/nest_asyncio-1.6.0-py3-none-any.whl", hash = "sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c", size = 5195, upload-time = "2024-01-21T14:25:17.223Z" }, +] + +[[package]] +name = "numpy" +version = "2.3.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/76/65/21b3bc86aac7b8f2862db1e808f1ea22b028e30a225a34a5ede9bf8678f2/numpy-2.3.5.tar.gz", hash = "sha256:784db1dcdab56bf0517743e746dfb0f885fc68d948aba86eeec2cba234bdf1c0", size = 20584950, upload-time = "2025-11-16T22:52:42.067Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/69/9cde09f36da4b5a505341180a3f2e6fadc352fd4d2b7096ce9778db83f1a/numpy-2.3.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d0f23b44f57077c1ede8c5f26b30f706498b4862d3ff0a7298b8411dd2f043ff", size = 16728251, upload-time = "2025-11-16T22:50:19.013Z" }, + { url = "https://files.pythonhosted.org/packages/79/fb/f505c95ceddd7027347b067689db71ca80bd5ecc926f913f1a23e65cf09b/numpy-2.3.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:aa5bc7c5d59d831d9773d1170acac7893ce3a5e130540605770ade83280e7188", size = 12254652, upload-time = "2025-11-16T22:50:21.487Z" }, + { url = "https://files.pythonhosted.org/packages/78/da/8c7738060ca9c31b30e9301ee0cf6c5ffdbf889d9593285a1cead337f9a5/numpy-2.3.5-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:ccc933afd4d20aad3c00bcef049cb40049f7f196e0397f1109dba6fed63267b0", size = 5083172, upload-time = "2025-11-16T22:50:24.562Z" }, + { url = "https://files.pythonhosted.org/packages/a4/b4/ee5bb2537fb9430fd2ef30a616c3672b991a4129bb1c7dcc42aa0abbe5d7/numpy-2.3.5-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:afaffc4393205524af9dfa400fa250143a6c3bc646c08c9f5e25a9f4b4d6a903", size = 6622990, upload-time = "2025-11-16T22:50:26.47Z" }, + { url = "https://files.pythonhosted.org/packages/95/03/dc0723a013c7d7c19de5ef29e932c3081df1c14ba582b8b86b5de9db7f0f/numpy-2.3.5-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c75442b2209b8470d6d5d8b1c25714270686f14c749028d2199c54e29f20b4d", size = 14248902, upload-time = "2025-11-16T22:50:28.861Z" }, + { url = "https://files.pythonhosted.org/packages/f5/10/ca162f45a102738958dcec8023062dad0cbc17d1ab99d68c4e4a6c45fb2b/numpy-2.3.5-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:11e06aa0af8c0f05104d56450d6093ee639e15f24ecf62d417329d06e522e017", size = 16597430, upload-time = "2025-11-16T22:50:31.56Z" }, + { url = "https://files.pythonhosted.org/packages/2a/51/c1e29be863588db58175175f057286900b4b3327a1351e706d5e0f8dd679/numpy-2.3.5-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ed89927b86296067b4f81f108a2271d8926467a8868e554eaf370fc27fa3ccaf", size = 16024551, upload-time = "2025-11-16T22:50:34.242Z" }, + { url = "https://files.pythonhosted.org/packages/83/68/8236589d4dbb87253d28259d04d9b814ec0ecce7cb1c7fed29729f4c3a78/numpy-2.3.5-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:51c55fe3451421f3a6ef9a9c1439e82101c57a2c9eab9feb196a62b1a10b58ce", size = 18533275, upload-time = "2025-11-16T22:50:37.651Z" }, + { url = "https://files.pythonhosted.org/packages/40/56/2932d75b6f13465239e3b7b7e511be27f1b8161ca2510854f0b6e521c395/numpy-2.3.5-cp313-cp313-win32.whl", hash = "sha256:1978155dd49972084bd6ef388d66ab70f0c323ddee6f693d539376498720fb7e", size = 6277637, upload-time = "2025-11-16T22:50:40.11Z" }, + { url = "https://files.pythonhosted.org/packages/0c/88/e2eaa6cffb115b85ed7c7c87775cb8bcf0816816bc98ca8dbfa2ee33fe6e/numpy-2.3.5-cp313-cp313-win_amd64.whl", hash = "sha256:00dc4e846108a382c5869e77c6ed514394bdeb3403461d25a829711041217d5b", size = 12779090, upload-time = "2025-11-16T22:50:42.503Z" }, + { url = "https://files.pythonhosted.org/packages/8f/88/3f41e13a44ebd4034ee17baa384acac29ba6a4fcc2aca95f6f08ca0447d1/numpy-2.3.5-cp313-cp313-win_arm64.whl", hash = "sha256:0472f11f6ec23a74a906a00b48a4dcf3849209696dff7c189714511268d103ae", size = 10194710, upload-time = "2025-11-16T22:50:44.971Z" }, + { url = "https://files.pythonhosted.org/packages/13/cb/71744144e13389d577f867f745b7df2d8489463654a918eea2eeb166dfc9/numpy-2.3.5-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:414802f3b97f3c1eef41e530aaba3b3c1620649871d8cb38c6eaff034c2e16bd", size = 16827292, upload-time = "2025-11-16T22:50:47.715Z" }, + { url = "https://files.pythonhosted.org/packages/71/80/ba9dc6f2a4398e7f42b708a7fdc841bb638d353be255655498edbf9a15a8/numpy-2.3.5-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5ee6609ac3604fa7780e30a03e5e241a7956f8e2fcfe547d51e3afa5247ac47f", size = 12378897, upload-time = "2025-11-16T22:50:51.327Z" }, + { url = "https://files.pythonhosted.org/packages/2e/6d/db2151b9f64264bcceccd51741aa39b50150de9b602d98ecfe7e0c4bff39/numpy-2.3.5-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:86d835afea1eaa143012a2d7a3f45a3adce2d7adc8b4961f0b362214d800846a", size = 5207391, upload-time = "2025-11-16T22:50:54.542Z" }, + { url = "https://files.pythonhosted.org/packages/80/ae/429bacace5ccad48a14c4ae5332f6aa8ab9f69524193511d60ccdfdc65fa/numpy-2.3.5-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:30bc11310e8153ca664b14c5f1b73e94bd0503681fcf136a163de856f3a50139", size = 6721275, upload-time = "2025-11-16T22:50:56.794Z" }, + { url = "https://files.pythonhosted.org/packages/74/5b/1919abf32d8722646a38cd527bc3771eb229a32724ee6ba340ead9b92249/numpy-2.3.5-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1062fde1dcf469571705945b0f221b73928f34a20c904ffb45db101907c3454e", size = 14306855, upload-time = "2025-11-16T22:50:59.208Z" }, + { url = "https://files.pythonhosted.org/packages/a5/87/6831980559434973bebc30cd9c1f21e541a0f2b0c280d43d3afd909b66d0/numpy-2.3.5-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ce581db493ea1a96c0556360ede6607496e8bf9b3a8efa66e06477267bc831e9", size = 16657359, upload-time = "2025-11-16T22:51:01.991Z" }, + { url = "https://files.pythonhosted.org/packages/dd/91/c797f544491ee99fd00495f12ebb7802c440c1915811d72ac5b4479a3356/numpy-2.3.5-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:cc8920d2ec5fa99875b670bb86ddeb21e295cb07aa331810d9e486e0b969d946", size = 16093374, upload-time = "2025-11-16T22:51:05.291Z" }, + { url = "https://files.pythonhosted.org/packages/74/a6/54da03253afcbe7a72785ec4da9c69fb7a17710141ff9ac5fcb2e32dbe64/numpy-2.3.5-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:9ee2197ef8c4f0dfe405d835f3b6a14f5fee7782b5de51ba06fb65fc9b36e9f1", size = 18594587, upload-time = "2025-11-16T22:51:08.585Z" }, + { url = "https://files.pythonhosted.org/packages/80/e9/aff53abbdd41b0ecca94285f325aff42357c6b5abc482a3fcb4994290b18/numpy-2.3.5-cp313-cp313t-win32.whl", hash = "sha256:70b37199913c1bd300ff6e2693316c6f869c7ee16378faf10e4f5e3275b299c3", size = 6405940, upload-time = "2025-11-16T22:51:11.541Z" }, + { url = "https://files.pythonhosted.org/packages/d5/81/50613fec9d4de5480de18d4f8ef59ad7e344d497edbef3cfd80f24f98461/numpy-2.3.5-cp313-cp313t-win_amd64.whl", hash = "sha256:b501b5fa195cc9e24fe102f21ec0a44dffc231d2af79950b451e0d99cea02234", size = 12920341, upload-time = "2025-11-16T22:51:14.312Z" }, + { url = "https://files.pythonhosted.org/packages/bb/ab/08fd63b9a74303947f34f0bd7c5903b9c5532c2d287bead5bdf4c556c486/numpy-2.3.5-cp313-cp313t-win_arm64.whl", hash = "sha256:a80afd79f45f3c4a7d341f13acbe058d1ca8ac017c165d3fa0d3de6bc1a079d7", size = 10262507, upload-time = "2025-11-16T22:51:16.846Z" }, + { url = "https://files.pythonhosted.org/packages/ba/97/1a914559c19e32d6b2e233cf9a6a114e67c856d35b1d6babca571a3e880f/numpy-2.3.5-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:bf06bc2af43fa8d32d30fae16ad965663e966b1a3202ed407b84c989c3221e82", size = 16735706, upload-time = "2025-11-16T22:51:19.558Z" }, + { url = "https://files.pythonhosted.org/packages/57/d4/51233b1c1b13ecd796311216ae417796b88b0616cfd8a33ae4536330748a/numpy-2.3.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:052e8c42e0c49d2575621c158934920524f6c5da05a1d3b9bab5d8e259e045f0", size = 12264507, upload-time = "2025-11-16T22:51:22.492Z" }, + { url = "https://files.pythonhosted.org/packages/45/98/2fe46c5c2675b8306d0b4a3ec3494273e93e1226a490f766e84298576956/numpy-2.3.5-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:1ed1ec893cff7040a02c8aa1c8611b94d395590d553f6b53629a4461dc7f7b63", size = 5093049, upload-time = "2025-11-16T22:51:25.171Z" }, + { url = "https://files.pythonhosted.org/packages/ce/0e/0698378989bb0ac5f1660c81c78ab1fe5476c1a521ca9ee9d0710ce54099/numpy-2.3.5-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:2dcd0808a421a482a080f89859a18beb0b3d1e905b81e617a188bd80422d62e9", size = 6626603, upload-time = "2025-11-16T22:51:27Z" }, + { url = "https://files.pythonhosted.org/packages/5e/a6/9ca0eecc489640615642a6cbc0ca9e10df70df38c4d43f5a928ff18d8827/numpy-2.3.5-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:727fd05b57df37dc0bcf1a27767a3d9a78cbbc92822445f32cc3436ba797337b", size = 14262696, upload-time = "2025-11-16T22:51:29.402Z" }, + { url = "https://files.pythonhosted.org/packages/c8/f6/07ec185b90ec9d7217a00eeeed7383b73d7e709dae2a9a021b051542a708/numpy-2.3.5-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fffe29a1ef00883599d1dc2c51aa2e5d80afe49523c261a74933df395c15c520", size = 16597350, upload-time = "2025-11-16T22:51:32.167Z" }, + { url = "https://files.pythonhosted.org/packages/75/37/164071d1dde6a1a84c9b8e5b414fa127981bad47adf3a6b7e23917e52190/numpy-2.3.5-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:8f7f0e05112916223d3f438f293abf0727e1181b5983f413dfa2fefc4098245c", size = 16040190, upload-time = "2025-11-16T22:51:35.403Z" }, + { url = "https://files.pythonhosted.org/packages/08/3c/f18b82a406b04859eb026d204e4e1773eb41c5be58410f41ffa511d114ae/numpy-2.3.5-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2e2eb32ddb9ccb817d620ac1d8dae7c3f641c1e5f55f531a33e8ab97960a75b8", size = 18536749, upload-time = "2025-11-16T22:51:39.698Z" }, + { url = "https://files.pythonhosted.org/packages/40/79/f82f572bf44cf0023a2fe8588768e23e1592585020d638999f15158609e1/numpy-2.3.5-cp314-cp314-win32.whl", hash = "sha256:66f85ce62c70b843bab1fb14a05d5737741e74e28c7b8b5a064de10142fad248", size = 6335432, upload-time = "2025-11-16T22:51:42.476Z" }, + { url = "https://files.pythonhosted.org/packages/a3/2e/235b4d96619931192c91660805e5e49242389742a7a82c27665021db690c/numpy-2.3.5-cp314-cp314-win_amd64.whl", hash = "sha256:e6a0bc88393d65807d751a614207b7129a310ca4fe76a74e5c7da5fa5671417e", size = 12919388, upload-time = "2025-11-16T22:51:45.275Z" }, + { url = "https://files.pythonhosted.org/packages/07/2b/29fd75ce45d22a39c61aad74f3d718e7ab67ccf839ca8b60866054eb15f8/numpy-2.3.5-cp314-cp314-win_arm64.whl", hash = "sha256:aeffcab3d4b43712bb7a60b65f6044d444e75e563ff6180af8f98dd4b905dfd2", size = 10476651, upload-time = "2025-11-16T22:51:47.749Z" }, + { url = "https://files.pythonhosted.org/packages/17/e1/f6a721234ebd4d87084cfa68d081bcba2f5cfe1974f7de4e0e8b9b2a2ba1/numpy-2.3.5-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:17531366a2e3a9e30762c000f2c43a9aaa05728712e25c11ce1dbe700c53ad41", size = 16834503, upload-time = "2025-11-16T22:51:50.443Z" }, + { url = "https://files.pythonhosted.org/packages/5c/1c/baf7ffdc3af9c356e1c135e57ab7cf8d247931b9554f55c467efe2c69eff/numpy-2.3.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:d21644de1b609825ede2f48be98dfde4656aefc713654eeee280e37cadc4e0ad", size = 12381612, upload-time = "2025-11-16T22:51:53.609Z" }, + { url = "https://files.pythonhosted.org/packages/74/91/f7f0295151407ddc9ba34e699013c32c3c91944f9b35fcf9281163dc1468/numpy-2.3.5-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:c804e3a5aba5460c73955c955bdbd5c08c354954e9270a2c1565f62e866bdc39", size = 5210042, upload-time = "2025-11-16T22:51:56.213Z" }, + { url = "https://files.pythonhosted.org/packages/2e/3b/78aebf345104ec50dd50a4d06ddeb46a9ff5261c33bcc58b1c4f12f85ec2/numpy-2.3.5-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:cc0a57f895b96ec78969c34f682c602bf8da1a0270b09bc65673df2e7638ec20", size = 6724502, upload-time = "2025-11-16T22:51:58.584Z" }, + { url = "https://files.pythonhosted.org/packages/02/c6/7c34b528740512e57ef1b7c8337ab0b4f0bddf34c723b8996c675bc2bc91/numpy-2.3.5-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:900218e456384ea676e24ea6a0417f030a3b07306d29d7ad843957b40a9d8d52", size = 14308962, upload-time = "2025-11-16T22:52:01.698Z" }, + { url = "https://files.pythonhosted.org/packages/80/35/09d433c5262bc32d725bafc619e095b6a6651caf94027a03da624146f655/numpy-2.3.5-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:09a1bea522b25109bf8e6f3027bd810f7c1085c64a0c7ce050c1676ad0ba010b", size = 16655054, upload-time = "2025-11-16T22:52:04.267Z" }, + { url = "https://files.pythonhosted.org/packages/7a/ab/6a7b259703c09a88804fa2430b43d6457b692378f6b74b356155283566ac/numpy-2.3.5-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:04822c00b5fd0323c8166d66c701dc31b7fbd252c100acd708c48f763968d6a3", size = 16091613, upload-time = "2025-11-16T22:52:08.651Z" }, + { url = "https://files.pythonhosted.org/packages/c2/88/330da2071e8771e60d1038166ff9d73f29da37b01ec3eb43cb1427464e10/numpy-2.3.5-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d6889ec4ec662a1a37eb4b4fb26b6100841804dac55bd9df579e326cdc146227", size = 18591147, upload-time = "2025-11-16T22:52:11.453Z" }, + { url = "https://files.pythonhosted.org/packages/51/41/851c4b4082402d9ea860c3626db5d5df47164a712cb23b54be028b184c1c/numpy-2.3.5-cp314-cp314t-win32.whl", hash = "sha256:93eebbcf1aafdf7e2ddd44c2923e2672e1010bddc014138b229e49725b4d6be5", size = 6479806, upload-time = "2025-11-16T22:52:14.641Z" }, + { url = "https://files.pythonhosted.org/packages/90/30/d48bde1dfd93332fa557cff1972fbc039e055a52021fbef4c2c4b1eefd17/numpy-2.3.5-cp314-cp314t-win_amd64.whl", hash = "sha256:c8a9958e88b65c3b27e22ca2a076311636850b612d6bbfb76e8d156aacde2aaf", size = 13105760, upload-time = "2025-11-16T22:52:17.975Z" }, + { url = "https://files.pythonhosted.org/packages/2d/fd/4b5eb0b3e888d86aee4d198c23acec7d214baaf17ea93c1adec94c9518b9/numpy-2.3.5-cp314-cp314t-win_arm64.whl", hash = "sha256:6203fdf9f3dc5bdaed7319ad8698e685c7a3be10819f41d32a0723e611733b42", size = 10545459, upload-time = "2025-11-16T22:52:20.55Z" }, +] + +[[package]] +name = "nvidia-cublas-cu12" +version = "12.9.1.4" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/6c/90d3f532f608a03a13c1d6c16c266ffa3828e8011b1549d3b61db2ad59f5/nvidia_cublas_cu12-12.9.1.4-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:7a950dae01add3b415a5a5cdc4ec818fb5858263e9cca59004bb99fdbbd3a5d6", size = 575006342, upload-time = "2025-06-05T20:04:16.902Z" }, + { url = "https://files.pythonhosted.org/packages/77/3c/aa88abe01f3be3d1f8f787d1d33dc83e76fec05945f9a28fbb41cfb99cd5/nvidia_cublas_cu12-12.9.1.4-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:453611eb21a7c1f2c2156ed9f3a45b691deda0440ec550860290dc901af5b4c2", size = 581242350, upload-time = "2025-06-05T20:04:51.979Z" }, +] + +[[package]] +name = "nvidia-cuda-cccl-cu12" +version = "12.9.27" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/7e/82e49956b046bdc506c789235c587d9b3ef58b8bc1782258c1e247229647/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d7898b38aa68beaa234d48f0868273702342a196d6e2e9d0ef058dca2390ebea", size = 3152245, upload-time = "2025-05-01T19:32:04.802Z" }, + { url = "https://files.pythonhosted.org/packages/18/2a/d4cd8506d2044e082f8cd921be57392e6a9b5ccd3ffdf050362430a3d5d5/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:37869e17ce2e1ecec6eddf1927cca0f8c34e64fd848d40453df559091e2d7117", size = 3152243, upload-time = "2025-05-01T19:32:13.955Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/78/351b5c8cdbd9a6b4fb0d6ee73fb176dcdc1b6b6ad47c2ffff5ae8ca4a1f7/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:791853b030602c6a11d08b5578edfb957cadea06e9d3b26adbf8d036135a4afe", size = 10077166, upload-time = "2025-06-05T20:01:01.385Z" }, + { url = "https://files.pythonhosted.org/packages/c1/2e/b84e32197e33f39907b455b83395a017e697c07a449a2b15fd07fc1c9981/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:096bcf334f13e1984ba36685ad4c1d6347db214de03dbb6eebb237b41d9d934f", size = 10814997, upload-time = "2025-06-05T20:01:10.168Z" }, +] + +[[package]] +name = "nvidia-cuda-nvcc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/48/b54a06168a2190572a312bfe4ce443687773eb61367ced31e064953dd2f7/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:5d6a0d32fdc7ea39917c20065614ae93add6f577d840233237ff08e9a38f58f0", size = 40546229, upload-time = "2025-06-05T20:01:53.357Z" }, + { url = "https://files.pythonhosted.org/packages/d6/5c/8cc072436787104bbbcbde1f76ab4a0d89e68f7cebc758dd2ad7913a43d0/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:44e1eca4d08926193a558d2434b1bf83d57b4d5743e0c431c0c83d51da1df62b", size = 39411138, upload-time = "2025-06-05T20:01:43.182Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/85/e4af82cc9202023862090bfca4ea827d533329e925c758f0cde964cb54b7/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:210cf05005a447e29214e9ce50851e83fc5f4358df8b453155d5e1918094dcb4", size = 89568129, upload-time = "2025-06-05T20:02:41.973Z" }, + { url = "https://files.pythonhosted.org/packages/64/eb/c2295044b8f3b3b08860e2f6a912b702fc92568a167259df5dddb78f325e/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:096d4de6bda726415dfaf3198d4f5c522b8e70139c97feef5cd2ca6d4cd9cead", size = 44528905, upload-time = "2025-06-05T20:02:29.754Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/e0/0279bd94539fda525e0c8538db29b72a5a8495b0c12173113471d28bce78/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:83469a846206f2a733db0c42e223589ab62fd2fabac4432d2f8802de4bded0a4", size = 3515012, upload-time = "2025-06-05T20:00:35.519Z" }, + { url = "https://files.pythonhosted.org/packages/bc/46/a92db19b8309581092a3add7e6fceb4c301a3fd233969856a8cbf042cd3c/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:25bba2dfb01d48a9b59ca474a1ac43c6ebf7011f1b0b8cc44f54eb6ac48a96c3", size = 3493179, upload-time = "2025-06-05T20:00:53.735Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.19.0.56" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/09/b8/277c51962ee46fa3e5b203ac5f76107c650f781d6891e681e28e6f3e9fe6/nvidia_cudnn_cu12-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:08caaf27fe556aca82a3ee3b5aa49a77e7de0cfcb7ff4e5c29da426387a8267e", size = 656910700, upload-time = "2026-02-03T20:40:25.508Z" }, + { url = "https://files.pythonhosted.org/packages/c5/41/65225d42fba06fb3dd3972485ea258e7dd07a40d6e01c95da6766ad87354/nvidia_cudnn_cu12-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:ac6ad90a075bb33a94f2b4cf4622eac13dd4dc65cf6dd9c7572a318516a36625", size = 657906812, upload-time = "2026-02-03T20:44:12.638Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.4.1.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/2b/76445b0af890da61b501fde30650a1a4bd910607261b209cccb5235d3daa/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1a28c9b12260a1aa7a8fd12f5ebd82d027963d635ba82ff39a1acfa7c4c0fbcf", size = 200822453, upload-time = "2025-06-05T20:05:27.889Z" }, + { url = "https://files.pythonhosted.org/packages/95/f4/61e6996dd20481ee834f57a8e9dca28b1869366a135e0d42e2aa8493bdd4/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c67884f2a7d276b4b80eb56a79322a95df592ae5e765cf1243693365ccab4e28", size = 200877592, upload-time = "2025-06-05T20:05:45.862Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.5.82" +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/03/99/686ff9bf3a82a531c62b1a5c614476e8dfa24a9d89067aeedf3592ee4538/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:62efa83e4ace59a4c734d052bb72158e888aa7b770e1a5f601682f16fe5b4fd2", size = 337869834, upload-time = "2025-06-05T20:06:53.125Z" }, + { url = "https://files.pythonhosted.org/packages/33/40/79b0c64d44d6c166c0964ec1d803d067f4a145cca23e23925fd351d0e642/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:15da72d1340d29b5b3cf3fd100e3cd53421dde36002eda6ed93811af63c40d88", size = 338117415, upload-time = "2025-06-05T20:07:16.809Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.10.65" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/6f/8710fbd17cdd1d0fc3fea7d36d5b65ce1933611c31e1861da330206b253a/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:221c73e7482dd93eda44e65ce567c031c07e2f93f6fa0ecd3ba876a195023e83", size = 366359408, upload-time = "2025-06-05T20:07:42.501Z" }, + { url = "https://files.pythonhosted.org/packages/12/46/b0fd4b04f86577921feb97d8e2cf028afe04f614d17fb5013de9282c9216/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:73060ce019ac064a057267c585bf1fd5a353734151f87472ff02b2c5c9984e78", size = 366465088, upload-time = "2025-06-05T20:08:20.413Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.29.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/28/cf/bcf8bb0c0030b1b9a345331f6281c37d2a8669758521eb93c382f6f87c8f/nvidia_nccl_cu12-2.29.3-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:6351b79dc7d2cc3d654ea1523616b9eeded71fe9c8da66b71eef9a5d1b2adad4", size = 289708535, upload-time = "2026-02-03T21:10:58.804Z" }, + { url = "https://files.pythonhosted.org/packages/31/5a/cac7d231f322b66caa16fd4b136ebc8e4b18b2805811c2d58dc47210cdea/nvidia_nccl_cu12-2.29.3-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:35ad42e7d5d722a83c36a3a478e281c20a5646383deaf1b9ed1a9ab7d61bed53", size = 289760316, upload-time = "2026-02-03T21:11:37.899Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e3f1171dbdc83c5932a45f0f4c99180a70de9bd2718c1ab77d14104f6d7147f9", size = 39748338, upload-time = "2025-06-05T20:10:25.613Z" }, + { url = "https://files.pythonhosted.org/packages/97/bc/2dcba8e70cf3115b400fef54f213bcd6715a3195eba000f8330f11e40c45/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:994a05ef08ef4b0b299829cde613a424382aff7efb08a7172c1fa616cc3af2ca", size = 39514880, upload-time = "2025-06-05T20:10:04.89Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.5.19" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-cccl-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/1b/06a698b584b67670e7f108381ecdf680d3a497c0567f5183644034f4212e/nvidia_nvshmem_cu12-3.5.19-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:333d91f68038539f20a591ef4d5cee25a85d440537b31809c4e42ae1700fccce", size = 152515031, upload-time = "2026-01-02T04:24:25.119Z" }, + { url = "https://files.pythonhosted.org/packages/64/b9/6ab941001c23cfb43499b5b0b7417b0bb4dfba3a29ffa2b06985422dad50/nvidia_nvshmem_cu12-3.5.19-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f29a23c5ee75461ec4342e17757f9e010369ada8eb0d441070049db787d9e51b", size = 152681200, upload-time = "2026-01-02T04:24:46.524Z" }, +] + +[[package]] +name = "onnxruntime" +version = "1.24.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flatbuffers" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "sympy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/77/7172ecfcbdabd92f338e694f38c325f6fab29a38fa0a8c3d1c85b9f4617c/onnxruntime-1.24.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:82e367770e8fba8a87ba9f4c04bb527e6d4d7204540f1390f202c27a3b759fb4", size = 17211381, upload-time = "2026-02-05T17:31:09.601Z" }, + { url = "https://files.pythonhosted.org/packages/79/5b/532a0d75b93bbd0da0e108b986097ebe164b84fbecfdf2ddbf7c8a3a2e83/onnxruntime-1.24.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1099f3629832580fedf415cfce2462a56cc9ca2b560d6300c24558e2ac049134", size = 15016000, upload-time = "2026-02-05T17:31:00.116Z" }, + { url = "https://files.pythonhosted.org/packages/f6/b5/40606c7bce0702975a077bc6668cd072cd77695fc5c0b3fcf59bdb1fe65e/onnxruntime-1.24.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6361dda4270f3939a625670bd67ae0982a49b7f923207450e28433abc9c3a83b", size = 17097637, upload-time = "2026-02-05T17:31:34.787Z" }, + { url = "https://files.pythonhosted.org/packages/d5/a0/9e8f7933796b466241b934585723c700d8fb6bde2de856e65335193d7c93/onnxruntime-1.24.1-cp313-cp313-win_amd64.whl", hash = "sha256:bd1e4aefe73b6b99aa303cd72562ab6de3cccb09088100f8ad1c974be13079c7", size = 12492467, upload-time = "2026-02-05T17:32:09.834Z" }, + { url = "https://files.pythonhosted.org/packages/fb/8a/ee07d86e35035f9fed42497af76435f5a613d4e8b6c537ea0f8ef9fa85da/onnxruntime-1.24.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:88a2b54dca00c90fca6303eedf13d49b5b4191d031372c2e85f5cffe4d86b79e", size = 15025407, upload-time = "2026-02-05T17:31:02.251Z" }, + { url = "https://files.pythonhosted.org/packages/fd/9e/ab3e1dda4b126313d240e1aaa87792ddb1f5ba6d03ca2f093a7c4af8c323/onnxruntime-1.24.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2dfbba602da840615ed5b431facda4b3a43b5d8276cf9e0dbf13d842df105838", size = 17099810, upload-time = "2026-02-05T17:31:37.537Z" }, + { url = "https://files.pythonhosted.org/packages/87/23/167d964414cee2af9c72af323b28d2c4cb35beed855c830a23f198265c79/onnxruntime-1.24.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:890c503ca187bc883c3aa72c53f2a604ec8e8444bdd1bf6ac243ec6d5e085202", size = 17214004, upload-time = "2026-02-05T17:31:11.917Z" }, + { url = "https://files.pythonhosted.org/packages/b4/24/6e5558fdd51027d6830cf411bc003ae12c64054826382e2fab89e99486a0/onnxruntime-1.24.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4da1b84b3bdeec543120df169e5e62a1445bf732fc2c7fb036c2f8a4090455e8", size = 15017034, upload-time = "2026-02-05T17:31:04.331Z" }, + { url = "https://files.pythonhosted.org/packages/91/d4/3cb1c9eaae1103265ed7eb00a3eaeb0d9ba51dc88edc398b7071c9553bed/onnxruntime-1.24.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:557753ec345efa227c6a65139f3d29c76330fcbd54cc10dd1b64232ebb939c13", size = 17097531, upload-time = "2026-02-05T17:31:40.303Z" }, + { url = "https://files.pythonhosted.org/packages/0f/da/4522b199c12db7c5b46aaf265ee0d741abe65ea912f6c0aaa2cc18a4654d/onnxruntime-1.24.1-cp314-cp314-win_amd64.whl", hash = "sha256:ea4942104805e868f3ddddfa1fbb58b04503a534d489ab2d1452bbfa345c78c2", size = 12795556, upload-time = "2026-02-05T17:32:11.886Z" }, + { url = "https://files.pythonhosted.org/packages/a1/53/3b8969417276b061ff04502ccdca9db4652d397abbeb06c9f6ae05cec9ca/onnxruntime-1.24.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ea8963a99e0f10489acdf00ef3383c3232b7e44aa497b063c63be140530d9f85", size = 15025434, upload-time = "2026-02-05T17:31:06.942Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a2/cfcf009eb38d90cc628c087b6506b3dfe1263387f3cbbf8d272af4fef957/onnxruntime-1.24.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:34488aa760fb5c2e6d06a7ca9241124eb914a6a06f70936a14c669d1b3df9598", size = 17099815, upload-time = "2026-02-05T17:31:43.092Z" }, +] + +[[package]] +name = "onnxruntime-gpu" +version = "1.24.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flatbuffers" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "sympy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/e7/4e19062e95d3701c0d32c228aa848ba4a1cc97651e53628d978dba8e1267/onnxruntime_gpu-1.24.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:db9acb0d0e59d93b4fa6b7fd44284ece4408d0acee73235d43ed343f8cee7ee5", size = 252629216, upload-time = "2026-02-05T17:24:24.604Z" }, + { url = "https://files.pythonhosted.org/packages/c4/82/223d7120d8a98b07c104ddecfb0cc2536188e566a4e9c2dee7572453f89c/onnxruntime_gpu-1.24.1-cp313-cp313-win_amd64.whl", hash = "sha256:59fdb40743f0722f3b859209f649ea160ca6bb42799e43f49b70a3ec5fc8c4ad", size = 207089285, upload-time = "2026-02-05T17:24:18.497Z" }, + { url = "https://files.pythonhosted.org/packages/ac/82/3159e57f09d7e6c8ad47d8ba8d5bd7494f383bc1071481cf38c9c8142bf9/onnxruntime_gpu-1.24.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:88ca04e1dffea2d4c3c79cf4de7f429e99059d085f21b3e775a8d36380cd5186", size = 252633977, upload-time = "2026-02-05T17:24:33.568Z" }, + { url = "https://files.pythonhosted.org/packages/c1/b4/51ad0ab878ff1456a831a0566b4db982a904e22f138e4b2c5f021bac517f/onnxruntime_gpu-1.24.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ced66900b1f48bddb62b5233925c3b56f8e008e2c34ebf8c060b20cae5842bcf", size = 252629039, upload-time = "2026-02-05T17:24:43.551Z" }, + { url = "https://files.pythonhosted.org/packages/9c/46/336d4e09a6af66532eedde5c8f03a73eaa91a046b408522259ab6a604363/onnxruntime_gpu-1.24.1-cp314-cp314-win_amd64.whl", hash = "sha256:129f6ae8b331a6507759597cd317b23e94aed6ead1da951f803c3328f2990b0c", size = 209487551, upload-time = "2026-02-05T17:24:26.373Z" }, + { url = "https://files.pythonhosted.org/packages/6a/94/a3b20276261f5e64dbd72bda656af988282cff01f18c2685953600e2f810/onnxruntime_gpu-1.24.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:de2cee7e12b0f4813c62f9a48df83fd01d066cc970400c832252cf3c155a6957", size = 252633096, upload-time = "2026-02-05T17:24:53.248Z" }, +] + +[[package]] +name = "opt-einsum" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/b9/2ac072041e899a52f20cf9510850ff58295003aa75525e58343591b0cbfb/opt_einsum-3.4.0.tar.gz", hash = "sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac", size = 63004, upload-time = "2024-09-26T14:33:24.483Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd", size = 71932, upload-time = "2024-09-26T14:33:23.039Z" }, +] + +[[package]] +name = "optax" +version = "0.2.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "jax" }, + { name = "jaxlib" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/de/f7/a63fc3d262d7a58d7d53050dea1408a63738739569af34f8f754cf181ab1/optax-0.2.7.tar.gz", hash = "sha256:8b6b2e5bd62bcc6c11f6172a1aff0d86da0eaeecbd5465b2b366b5d3d64f6efc", size = 297524, upload-time = "2026-02-05T20:49:28.749Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/1e/94ad43e06887244b4d25f58b689122270ba3c129d3448052958eecf7518a/optax-0.2.7-py3-none-any.whl", hash = "sha256:241f2dfa104eab4fec2e16e7919f88df24a3da1481f95e264b3db396b30d4ff6", size = 399395, upload-time = "2026-02-05T20:49:26.883Z" }, +] + +[[package]] +name = "orbax-checkpoint" +version = "0.11.32" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "aiofiles" }, + { name = "etils", extra = ["epath", "epy"] }, + { name = "humanize" }, + { name = "jax" }, + { name = "msgpack" }, + { name = "nest-asyncio" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "psutil" }, + { name = "pyyaml" }, + { name = "simplejson" }, + { name = "tensorstore" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6c/5f/1733e1143696319f311bc4de48da2e306a1f62f0925f9fe9d797b8ba8abe/orbax_checkpoint-0.11.32.tar.gz", hash = "sha256:523dcf61e93c7187c6b80fd50f3177114c0b957ea62cbb5c869c0b3e3d1a7dfc", size = 431601, upload-time = "2026-01-20T16:46:06.307Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/17/aae3144258f30920741ec91dbff0ff54665e572da50e6445ef437e08ec32/orbax_checkpoint-0.11.32-py3-none-any.whl", hash = "sha256:f0bfe9f9b1ce2c32c8f5dfab63393e51de525d41352abc17c7e21f9cc731d7a9", size = 634424, upload-time = "2026-01-20T16:46:04.382Z" }, +] + +[[package]] +name = "orbax-export" +version = "0.0.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "dataclasses-json" }, + { name = "etils" }, + { name = "jax" }, + { name = "jaxlib" }, + { name = "jaxtyping" }, + { name = "numpy" }, + { name = "orbax-checkpoint" }, + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1c/c8/ed7ac3c3c687bf129d7469b016c2b3d8777379f4ea453474e50ee41ce5cb/orbax_export-0.0.8.tar.gz", hash = "sha256:544eef564e2a6f17cd11b1167febe348b7b7cf56d9575de994a33d5613dd568a", size = 124980, upload-time = "2025-09-17T15:41:14.264Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/a9/3a755a58c8b6a36fe7e9e66bb6b93967ff49cdbc77cca8eacb2cf66435e9/orbax_export-0.0.8-py3-none-any.whl", hash = "sha256:f8037e1666ad28411cdb08d0668a2737b1281a32902c623ceda12109a089bc36", size = 180487, upload-time = "2025-09-17T15:41:12.928Z" }, +] + +[[package]] +name = "packaging" +version = "26.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/65/ee/299d360cdc32edc7d2cf530f3accf79c4fca01e96ffc950d8a52213bd8e4/packaging-26.0.tar.gz", hash = "sha256:00243ae351a257117b6a241061796684b084ed1c516a08c48a3f7e147a9d80b4", size = 143416, upload-time = "2026-01-21T20:50:39.064Z" } +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 = "pathspec" +version = "1.0.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fa/36/e27608899f9b8d4dff0617b2d9ab17ca5608956ca44461ac14ac48b44015/pathspec-1.0.4.tar.gz", hash = "sha256:0210e2ae8a21a9137c0d470578cb0e595af87edaa6ebf12ff176f14a02e0e645", size = 131200, upload-time = "2026-01-27T03:59:46.938Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ef/3c/2c197d226f9ea224a9ab8d197933f9da0ae0aac5b6e0f884e2b8d9c8e9f7/pathspec-1.0.4-py3-none-any.whl", hash = "sha256:fb6ae2fd4e7c921a165808a552060e722767cfa526f99ca5156ed2ce45a5c723", size = 55206, upload-time = "2026-01-27T03:59:45.137Z" }, +] + +[[package]] +name = "pillow" +version = "12.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/02/d52c733a2452ef1ffcc123b68e6606d07276b0e358db70eabad7e40042b7/pillow-12.1.0.tar.gz", hash = "sha256:5c5ae0a06e9ea030ab786b0251b32c7e4ce10e58d983c0d5c56029455180b5b9", size = 46977283, upload-time = "2026-01-02T09:13:29.892Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dd/c7/2530a4aa28248623e9d7f27316b42e27c32ec410f695929696f2e0e4a778/pillow-12.1.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:7b5dd7cbae20285cdb597b10eb5a2c13aa9de6cde9bb64a3c1317427b1db1ae1", size = 4062543, upload-time = "2026-01-02T09:11:31.566Z" }, + { url = "https://files.pythonhosted.org/packages/8f/1f/40b8eae823dc1519b87d53c30ed9ef085506b05281d313031755c1705f73/pillow-12.1.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:29a4cef9cb672363926f0470afc516dbf7305a14d8c54f7abbb5c199cd8f8179", size = 4138373, upload-time = "2026-01-02T09:11:33.367Z" }, + { url = "https://files.pythonhosted.org/packages/d4/77/6fa60634cf06e52139fd0e89e5bbf055e8166c691c42fb162818b7fda31d/pillow-12.1.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:681088909d7e8fa9e31b9799aaa59ba5234c58e5e4f1951b4c4d1082a2e980e0", size = 3601241, upload-time = "2026-01-02T09:11:35.011Z" }, + { url = "https://files.pythonhosted.org/packages/4f/bf/28ab865de622e14b747f0cd7877510848252d950e43002e224fb1c9ababf/pillow-12.1.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:983976c2ab753166dc66d36af6e8ec15bb511e4a25856e2227e5f7e00a160587", size = 5262410, upload-time = "2026-01-02T09:11:36.682Z" }, + { url = "https://files.pythonhosted.org/packages/1c/34/583420a1b55e715937a85bd48c5c0991598247a1fd2eb5423188e765ea02/pillow-12.1.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:db44d5c160a90df2d24a24760bbd37607d53da0b34fb546c4c232af7192298ac", size = 4657312, upload-time = "2026-01-02T09:11:38.535Z" }, + { url = "https://files.pythonhosted.org/packages/1d/fd/f5a0896839762885b3376ff04878f86ab2b097c2f9a9cdccf4eda8ba8dc0/pillow-12.1.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6b7a9d1db5dad90e2991645874f708e87d9a3c370c243c2d7684d28f7e133e6b", size = 6232605, upload-time = "2026-01-02T09:11:40.602Z" }, + { url = "https://files.pythonhosted.org/packages/98/aa/938a09d127ac1e70e6ed467bd03834350b33ef646b31edb7452d5de43792/pillow-12.1.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6258f3260986990ba2fa8a874f8b6e808cf5abb51a94015ca3dc3c68aa4f30ea", size = 8041617, upload-time = "2026-01-02T09:11:42.721Z" }, + { url = "https://files.pythonhosted.org/packages/17/e8/538b24cb426ac0186e03f80f78bc8dc7246c667f58b540bdd57c71c9f79d/pillow-12.1.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e115c15e3bc727b1ca3e641a909f77f8ca72a64fff150f666fcc85e57701c26c", size = 6346509, upload-time = "2026-01-02T09:11:44.955Z" }, + { url = "https://files.pythonhosted.org/packages/01/9a/632e58ec89a32738cabfd9ec418f0e9898a2b4719afc581f07c04a05e3c9/pillow-12.1.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6741e6f3074a35e47c77b23a4e4f2d90db3ed905cb1c5e6e0d49bff2045632bc", size = 7038117, upload-time = "2026-01-02T09:11:46.736Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a2/d40308cf86eada842ca1f3ffa45d0ca0df7e4ab33c83f81e73f5eaed136d/pillow-12.1.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:935b9d1aed48fcfb3f838caac506f38e29621b44ccc4f8a64d575cb1b2a88644", size = 6460151, upload-time = "2026-01-02T09:11:48.625Z" }, + { url = "https://files.pythonhosted.org/packages/f1/88/f5b058ad6453a085c5266660a1417bdad590199da1b32fb4efcff9d33b05/pillow-12.1.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5fee4c04aad8932da9f8f710af2c1a15a83582cfb884152a9caa79d4efcdbf9c", size = 7164534, upload-time = "2026-01-02T09:11:50.445Z" }, + { url = "https://files.pythonhosted.org/packages/19/ce/c17334caea1db789163b5d855a5735e47995b0b5dc8745e9a3605d5f24c0/pillow-12.1.0-cp313-cp313-win32.whl", hash = "sha256:a786bf667724d84aa29b5db1c61b7bfdde380202aaca12c3461afd6b71743171", size = 6332551, upload-time = "2026-01-02T09:11:52.234Z" }, + { url = "https://files.pythonhosted.org/packages/e5/07/74a9d941fa45c90a0d9465098fe1ec85de3e2afbdc15cc4766622d516056/pillow-12.1.0-cp313-cp313-win_amd64.whl", hash = "sha256:461f9dfdafa394c59cd6d818bdfdbab4028b83b02caadaff0ffd433faf4c9a7a", size = 7040087, upload-time = "2026-01-02T09:11:54.822Z" }, + { url = "https://files.pythonhosted.org/packages/88/09/c99950c075a0e9053d8e880595926302575bc742b1b47fe1bbcc8d388d50/pillow-12.1.0-cp313-cp313-win_arm64.whl", hash = "sha256:9212d6b86917a2300669511ed094a9406888362e085f2431a7da985a6b124f45", size = 2452470, upload-time = "2026-01-02T09:11:56.522Z" }, + { url = "https://files.pythonhosted.org/packages/b5/ba/970b7d85ba01f348dee4d65412476321d40ee04dcb51cd3735b9dc94eb58/pillow-12.1.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:00162e9ca6d22b7c3ee8e61faa3c3253cd19b6a37f126cad04f2f88b306f557d", size = 5264816, upload-time = "2026-01-02T09:11:58.227Z" }, + { url = "https://files.pythonhosted.org/packages/10/60/650f2fb55fdba7a510d836202aa52f0baac633e50ab1cf18415d332188fb/pillow-12.1.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:7d6daa89a00b58c37cb1747ec9fb7ac3bc5ffd5949f5888657dfddde6d1312e0", size = 4660472, upload-time = "2026-01-02T09:12:00.798Z" }, + { url = "https://files.pythonhosted.org/packages/2b/c0/5273a99478956a099d533c4f46cbaa19fd69d606624f4334b85e50987a08/pillow-12.1.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e2479c7f02f9d505682dc47df8c0ea1fc5e264c4d1629a5d63fe3e2334b89554", size = 6268974, upload-time = "2026-01-02T09:12:02.572Z" }, + { url = "https://files.pythonhosted.org/packages/b4/26/0bf714bc2e73d5267887d47931d53c4ceeceea6978148ed2ab2a4e6463c4/pillow-12.1.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f188d580bd870cda1e15183790d1cc2fa78f666e76077d103edf048eed9c356e", size = 8073070, upload-time = "2026-01-02T09:12:04.75Z" }, + { url = "https://files.pythonhosted.org/packages/43/cf/1ea826200de111a9d65724c54f927f3111dc5ae297f294b370a670c17786/pillow-12.1.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0fde7ec5538ab5095cc02df38ee99b0443ff0e1c847a045554cf5f9af1f4aa82", size = 6380176, upload-time = "2026-01-02T09:12:06.626Z" }, + { url = "https://files.pythonhosted.org/packages/03/e0/7938dd2b2013373fd85d96e0f38d62b7a5a262af21ac274250c7ca7847c9/pillow-12.1.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0ed07dca4a8464bada6139ab38f5382f83e5f111698caf3191cb8dbf27d908b4", size = 7067061, upload-time = "2026-01-02T09:12:08.624Z" }, + { url = "https://files.pythonhosted.org/packages/86/ad/a2aa97d37272a929a98437a8c0ac37b3cf012f4f8721e1bd5154699b2518/pillow-12.1.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:f45bd71d1fa5e5749587613037b172e0b3b23159d1c00ef2fc920da6f470e6f0", size = 6491824, upload-time = "2026-01-02T09:12:10.488Z" }, + { url = "https://files.pythonhosted.org/packages/a4/44/80e46611b288d51b115826f136fb3465653c28f491068a72d3da49b54cd4/pillow-12.1.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:277518bf4fe74aa91489e1b20577473b19ee70fb97c374aa50830b279f25841b", size = 7190911, upload-time = "2026-01-02T09:12:12.772Z" }, + { url = "https://files.pythonhosted.org/packages/86/77/eacc62356b4cf81abe99ff9dbc7402750044aed02cfd6a503f7c6fc11f3e/pillow-12.1.0-cp313-cp313t-win32.whl", hash = "sha256:7315f9137087c4e0ee73a761b163fc9aa3b19f5f606a7fc08d83fd3e4379af65", size = 6336445, upload-time = "2026-01-02T09:12:14.775Z" }, + { url = "https://files.pythonhosted.org/packages/e7/3c/57d81d0b74d218706dafccb87a87ea44262c43eef98eb3b164fd000e0491/pillow-12.1.0-cp313-cp313t-win_amd64.whl", hash = "sha256:0ddedfaa8b5f0b4ffbc2fa87b556dc59f6bb4ecb14a53b33f9189713ae8053c0", size = 7045354, upload-time = "2026-01-02T09:12:16.599Z" }, + { url = "https://files.pythonhosted.org/packages/ac/82/8b9b97bba2e3576a340f93b044a3a3a09841170ab4c1eb0d5c93469fd32f/pillow-12.1.0-cp313-cp313t-win_arm64.whl", hash = "sha256:80941e6d573197a0c28f394753de529bb436b1ca990ed6e765cf42426abc39f8", size = 2454547, upload-time = "2026-01-02T09:12:18.704Z" }, + { url = "https://files.pythonhosted.org/packages/8c/87/bdf971d8bbcf80a348cc3bacfcb239f5882100fe80534b0ce67a784181d8/pillow-12.1.0-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:5cb7bc1966d031aec37ddb9dcf15c2da5b2e9f7cc3ca7c54473a20a927e1eb91", size = 4062533, upload-time = "2026-01-02T09:12:20.791Z" }, + { url = "https://files.pythonhosted.org/packages/ff/4f/5eb37a681c68d605eb7034c004875c81f86ec9ef51f5be4a63eadd58859a/pillow-12.1.0-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:97e9993d5ed946aba26baf9c1e8cf18adbab584b99f452ee72f7ee8acb882796", size = 4138546, upload-time = "2026-01-02T09:12:23.664Z" }, + { url = "https://files.pythonhosted.org/packages/11/6d/19a95acb2edbace40dcd582d077b991646b7083c41b98da4ed7555b59733/pillow-12.1.0-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:414b9a78e14ffeb98128863314e62c3f24b8a86081066625700b7985b3f529bd", size = 3601163, upload-time = "2026-01-02T09:12:26.338Z" }, + { url = "https://files.pythonhosted.org/packages/fc/36/2b8138e51cb42e4cc39c3297713455548be855a50558c3ac2beebdc251dd/pillow-12.1.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:e6bdb408f7c9dd2a5ff2b14a3b0bb6d4deb29fb9961e6eb3ae2031ae9a5cec13", size = 5266086, upload-time = "2026-01-02T09:12:28.782Z" }, + { url = "https://files.pythonhosted.org/packages/53/4b/649056e4d22e1caa90816bf99cef0884aed607ed38075bd75f091a607a38/pillow-12.1.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:3413c2ae377550f5487991d444428f1a8ae92784aac79caa8b1e3b89b175f77e", size = 4657344, upload-time = "2026-01-02T09:12:31.117Z" }, + { url = "https://files.pythonhosted.org/packages/6c/6b/c5742cea0f1ade0cd61485dc3d81f05261fc2276f537fbdc00802de56779/pillow-12.1.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e5dcbe95016e88437ecf33544ba5db21ef1b8dd6e1b434a2cb2a3d605299e643", size = 6232114, upload-time = "2026-01-02T09:12:32.936Z" }, + { url = "https://files.pythonhosted.org/packages/bf/8f/9f521268ce22d63991601aafd3d48d5ff7280a246a1ef62d626d67b44064/pillow-12.1.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d0a7735df32ccbcc98b98a1ac785cc4b19b580be1bdf0aeb5c03223220ea09d5", size = 8042708, upload-time = "2026-01-02T09:12:34.78Z" }, + { url = "https://files.pythonhosted.org/packages/1a/eb/257f38542893f021502a1bbe0c2e883c90b5cff26cc33b1584a841a06d30/pillow-12.1.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0c27407a2d1b96774cbc4a7594129cc027339fd800cd081e44497722ea1179de", size = 6347762, upload-time = "2026-01-02T09:12:36.748Z" }, + { url = "https://files.pythonhosted.org/packages/c4/5a/8ba375025701c09b309e8d5163c5a4ce0102fa86bbf8800eb0d7ac87bc51/pillow-12.1.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:15c794d74303828eaa957ff8070846d0efe8c630901a1c753fdc63850e19ecd9", size = 7039265, upload-time = "2026-01-02T09:12:39.082Z" }, + { url = "https://files.pythonhosted.org/packages/cf/dc/cf5e4cdb3db533f539e88a7bbf9f190c64ab8a08a9bc7a4ccf55067872e4/pillow-12.1.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c990547452ee2800d8506c4150280757f88532f3de2a58e3022e9b179107862a", size = 6462341, upload-time = "2026-01-02T09:12:40.946Z" }, + { url = "https://files.pythonhosted.org/packages/d0/47/0291a25ac9550677e22eda48510cfc4fa4b2ef0396448b7fbdc0a6946309/pillow-12.1.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:b63e13dd27da389ed9475b3d28510f0f954bca0041e8e551b2a4eb1eab56a39a", size = 7165395, upload-time = "2026-01-02T09:12:42.706Z" }, + { url = "https://files.pythonhosted.org/packages/4f/4c/e005a59393ec4d9416be06e6b45820403bb946a778e39ecec62f5b2b991e/pillow-12.1.0-cp314-cp314-win32.whl", hash = "sha256:1a949604f73eb07a8adab38c4fe50791f9919344398bdc8ac6b307f755fc7030", size = 6431413, upload-time = "2026-01-02T09:12:44.944Z" }, + { url = "https://files.pythonhosted.org/packages/1c/af/f23697f587ac5f9095d67e31b81c95c0249cd461a9798a061ed6709b09b5/pillow-12.1.0-cp314-cp314-win_amd64.whl", hash = "sha256:4f9f6a650743f0ddee5593ac9e954ba1bdbc5e150bc066586d4f26127853ab94", size = 7176779, upload-time = "2026-01-02T09:12:46.727Z" }, + { url = "https://files.pythonhosted.org/packages/b3/36/6a51abf8599232f3e9afbd16d52829376a68909fe14efe29084445db4b73/pillow-12.1.0-cp314-cp314-win_arm64.whl", hash = "sha256:808b99604f7873c800c4840f55ff389936ef1948e4e87645eaf3fccbc8477ac4", size = 2543105, upload-time = "2026-01-02T09:12:49.243Z" }, + { url = "https://files.pythonhosted.org/packages/82/54/2e1dd20c8749ff225080d6ba465a0cab4387f5db0d1c5fb1439e2d99923f/pillow-12.1.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:bc11908616c8a283cf7d664f77411a5ed2a02009b0097ff8abbba5e79128ccf2", size = 5268571, upload-time = "2026-01-02T09:12:51.11Z" }, + { url = "https://files.pythonhosted.org/packages/57/61/571163a5ef86ec0cf30d265ac2a70ae6fc9e28413d1dc94fa37fae6bda89/pillow-12.1.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:896866d2d436563fa2a43a9d72f417874f16b5545955c54a64941e87c1376c61", size = 4660426, upload-time = "2026-01-02T09:12:52.865Z" }, + { url = "https://files.pythonhosted.org/packages/5e/e1/53ee5163f794aef1bf84243f755ee6897a92c708505350dd1923f4afec48/pillow-12.1.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8e178e3e99d3c0ea8fc64b88447f7cac8ccf058af422a6cedc690d0eadd98c51", size = 6269908, upload-time = "2026-01-02T09:12:54.884Z" }, + { url = "https://files.pythonhosted.org/packages/bc/0b/b4b4106ff0ee1afa1dc599fde6ab230417f800279745124f6c50bcffed8e/pillow-12.1.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:079af2fb0c599c2ec144ba2c02766d1b55498e373b3ac64687e43849fbbef5bc", size = 8074733, upload-time = "2026-01-02T09:12:56.802Z" }, + { url = "https://files.pythonhosted.org/packages/19/9f/80b411cbac4a732439e629a26ad3ef11907a8c7fc5377b7602f04f6fe4e7/pillow-12.1.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bdec5e43377761c5dbca620efb69a77f6855c5a379e32ac5b158f54c84212b14", size = 6381431, upload-time = "2026-01-02T09:12:58.823Z" }, + { url = "https://files.pythonhosted.org/packages/8f/b7/d65c45db463b66ecb6abc17c6ba6917a911202a07662247e1355ce1789e7/pillow-12.1.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:565c986f4b45c020f5421a4cea13ef294dde9509a8577f29b2fc5edc7587fff8", size = 7068529, upload-time = "2026-01-02T09:13:00.885Z" }, + { url = "https://files.pythonhosted.org/packages/50/96/dfd4cd726b4a45ae6e3c669fc9e49deb2241312605d33aba50499e9d9bd1/pillow-12.1.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:43aca0a55ce1eefc0aefa6253661cb54571857b1a7b2964bd8a1e3ef4b729924", size = 6492981, upload-time = "2026-01-02T09:13:03.314Z" }, + { url = "https://files.pythonhosted.org/packages/4d/1c/b5dc52cf713ae46033359c5ca920444f18a6359ce1020dd3e9c553ea5bc6/pillow-12.1.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:0deedf2ea233722476b3a81e8cdfbad786f7adbed5d848469fa59fe52396e4ef", size = 7191878, upload-time = "2026-01-02T09:13:05.276Z" }, + { url = "https://files.pythonhosted.org/packages/53/26/c4188248bd5edaf543864fe4834aebe9c9cb4968b6f573ce014cc42d0720/pillow-12.1.0-cp314-cp314t-win32.whl", hash = "sha256:b17fbdbe01c196e7e159aacb889e091f28e61020a8abeac07b68079b6e626988", size = 6438703, upload-time = "2026-01-02T09:13:07.491Z" }, + { url = "https://files.pythonhosted.org/packages/b8/0e/69ed296de8ea05cb03ee139cee600f424ca166e632567b2d66727f08c7ed/pillow-12.1.0-cp314-cp314t-win_amd64.whl", hash = "sha256:27b9baecb428899db6c0de572d6d305cfaf38ca1596b5c0542a5182e3e74e8c6", size = 7182927, upload-time = "2026-01-02T09:13:09.841Z" }, + { url = "https://files.pythonhosted.org/packages/fc/f5/68334c015eed9b5cff77814258717dec591ded209ab5b6fb70e2ae873d1d/pillow-12.1.0-cp314-cp314t-win_arm64.whl", hash = "sha256:f61333d817698bdcdd0f9d7793e365ac3d2a21c1f1eb02b32ad6aefb8d8ea831", size = 2545104, upload-time = "2026-01-02T09:13:12.068Z" }, +] + +[[package]] +name = "platformdirs" +version = "4.5.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cf/86/0248f086a84f01b37aaec0fa567b397df1a119f73c16f6c7a9aac73ea309/platformdirs-4.5.1.tar.gz", hash = "sha256:61d5cdcc6065745cdd94f0f878977f8de9437be93de97c1c12f853c9c0cdcbda", size = 21715, upload-time = "2025-12-05T13:52:58.638Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/28/3bfe2fa5a7b9c46fe7e13c97bda14c895fb10fa2ebf1d0abb90e0cea7ee1/platformdirs-4.5.1-py3-none-any.whl", hash = "sha256:d03afa3963c806a9bed9d5125c8f4cb2fdaf74a55ab60e5d59b3fde758104d31", size = 18731, upload-time = "2025-12-05T13:52:56.823Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +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 = "propcache" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9e/da/e9fc233cf63743258bff22b3dfa7ea5baef7b5bc324af47a0ad89b8ffc6f/propcache-0.4.1.tar.gz", hash = "sha256:f48107a8c637e80362555f37ecf49abe20370e557cc4ab374f04ec4423c97c3d", size = 46442, upload-time = "2025-10-08T19:49:02.291Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bf/df/6d9c1b6ac12b003837dde8a10231a7344512186e87b36e855bef32241942/propcache-0.4.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:43eedf29202c08550aac1d14e0ee619b0430aaef78f85864c1a892294fbc28cf", size = 77750, upload-time = "2025-10-08T19:47:07.648Z" }, + { url = "https://files.pythonhosted.org/packages/8b/e8/677a0025e8a2acf07d3418a2e7ba529c9c33caf09d3c1f25513023c1db56/propcache-0.4.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d62cdfcfd89ccb8de04e0eda998535c406bf5e060ffd56be6c586cbcc05b3311", size = 44780, upload-time = "2025-10-08T19:47:08.851Z" }, + { url = "https://files.pythonhosted.org/packages/89/a4/92380f7ca60f99ebae761936bc48a72a639e8a47b29050615eef757cb2a7/propcache-0.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:cae65ad55793da34db5f54e4029b89d3b9b9490d8abe1b4c7ab5d4b8ec7ebf74", size = 46308, upload-time = "2025-10-08T19:47:09.982Z" }, + { url = "https://files.pythonhosted.org/packages/2d/48/c5ac64dee5262044348d1d78a5f85dd1a57464a60d30daee946699963eb3/propcache-0.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:333ddb9031d2704a301ee3e506dc46b1fe5f294ec198ed6435ad5b6a085facfe", size = 208182, upload-time = "2025-10-08T19:47:11.319Z" }, + { url = "https://files.pythonhosted.org/packages/c6/0c/cd762dd011a9287389a6a3eb43aa30207bde253610cca06824aeabfe9653/propcache-0.4.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:fd0858c20f078a32cf55f7e81473d96dcf3b93fd2ccdb3d40fdf54b8573df3af", size = 211215, upload-time = "2025-10-08T19:47:13.146Z" }, + { url = "https://files.pythonhosted.org/packages/30/3e/49861e90233ba36890ae0ca4c660e95df565b2cd15d4a68556ab5865974e/propcache-0.4.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:678ae89ebc632c5c204c794f8dab2837c5f159aeb59e6ed0539500400577298c", size = 218112, upload-time = "2025-10-08T19:47:14.913Z" }, + { url = "https://files.pythonhosted.org/packages/f1/8b/544bc867e24e1bd48f3118cecd3b05c694e160a168478fa28770f22fd094/propcache-0.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d472aeb4fbf9865e0c6d622d7f4d54a4e101a89715d8904282bb5f9a2f476c3f", size = 204442, upload-time = "2025-10-08T19:47:16.277Z" }, + { url = "https://files.pythonhosted.org/packages/50/a6/4282772fd016a76d3e5c0df58380a5ea64900afd836cec2c2f662d1b9bb3/propcache-0.4.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4d3df5fa7e36b3225954fba85589da77a0fe6a53e3976de39caf04a0db4c36f1", size = 199398, upload-time = "2025-10-08T19:47:17.962Z" }, + { url = "https://files.pythonhosted.org/packages/3e/ec/d8a7cd406ee1ddb705db2139f8a10a8a427100347bd698e7014351c7af09/propcache-0.4.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:ee17f18d2498f2673e432faaa71698032b0127ebf23ae5974eeaf806c279df24", size = 196920, upload-time = "2025-10-08T19:47:19.355Z" }, + { url = "https://files.pythonhosted.org/packages/f6/6c/f38ab64af3764f431e359f8baf9e0a21013e24329e8b85d2da32e8ed07ca/propcache-0.4.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:580e97762b950f993ae618e167e7be9256b8353c2dcd8b99ec100eb50f5286aa", size = 203748, upload-time = "2025-10-08T19:47:21.338Z" }, + { url = "https://files.pythonhosted.org/packages/d6/e3/fa846bd70f6534d647886621388f0a265254d30e3ce47e5c8e6e27dbf153/propcache-0.4.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:501d20b891688eb8e7aa903021f0b72d5a55db40ffaab27edefd1027caaafa61", size = 205877, upload-time = "2025-10-08T19:47:23.059Z" }, + { url = "https://files.pythonhosted.org/packages/e2/39/8163fc6f3133fea7b5f2827e8eba2029a0277ab2c5beee6c1db7b10fc23d/propcache-0.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9a0bd56e5b100aef69bd8562b74b46254e7c8812918d3baa700c8a8009b0af66", size = 199437, upload-time = "2025-10-08T19:47:24.445Z" }, + { url = "https://files.pythonhosted.org/packages/93/89/caa9089970ca49c7c01662bd0eeedfe85494e863e8043565aeb6472ce8fe/propcache-0.4.1-cp313-cp313-win32.whl", hash = "sha256:bcc9aaa5d80322bc2fb24bb7accb4a30f81e90ab8d6ba187aec0744bc302ad81", size = 37586, upload-time = "2025-10-08T19:47:25.736Z" }, + { url = "https://files.pythonhosted.org/packages/f5/ab/f76ec3c3627c883215b5c8080debb4394ef5a7a29be811f786415fc1e6fd/propcache-0.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:381914df18634f5494334d201e98245c0596067504b9372d8cf93f4bb23e025e", size = 40790, upload-time = "2025-10-08T19:47:26.847Z" }, + { url = "https://files.pythonhosted.org/packages/59/1b/e71ae98235f8e2ba5004d8cb19765a74877abf189bc53fc0c80d799e56c3/propcache-0.4.1-cp313-cp313-win_arm64.whl", hash = "sha256:8873eb4460fd55333ea49b7d189749ecf6e55bf85080f11b1c4530ed3034cba1", size = 37158, upload-time = "2025-10-08T19:47:27.961Z" }, + { url = "https://files.pythonhosted.org/packages/83/ce/a31bbdfc24ee0dcbba458c8175ed26089cf109a55bbe7b7640ed2470cfe9/propcache-0.4.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:92d1935ee1f8d7442da9c0c4fa7ac20d07e94064184811b685f5c4fada64553b", size = 81451, upload-time = "2025-10-08T19:47:29.445Z" }, + { url = "https://files.pythonhosted.org/packages/25/9c/442a45a470a68456e710d96cacd3573ef26a1d0a60067e6a7d5e655621ed/propcache-0.4.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:473c61b39e1460d386479b9b2f337da492042447c9b685f28be4f74d3529e566", size = 46374, upload-time = "2025-10-08T19:47:30.579Z" }, + { url = "https://files.pythonhosted.org/packages/f4/bf/b1d5e21dbc3b2e889ea4327044fb16312a736d97640fb8b6aa3f9c7b3b65/propcache-0.4.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:c0ef0aaafc66fbd87842a3fe3902fd889825646bc21149eafe47be6072725835", size = 48396, upload-time = "2025-10-08T19:47:31.79Z" }, + { url = "https://files.pythonhosted.org/packages/f4/04/5b4c54a103d480e978d3c8a76073502b18db0c4bc17ab91b3cb5092ad949/propcache-0.4.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f95393b4d66bfae908c3ca8d169d5f79cd65636ae15b5e7a4f6e67af675adb0e", size = 275950, upload-time = "2025-10-08T19:47:33.481Z" }, + { url = "https://files.pythonhosted.org/packages/b4/c1/86f846827fb969c4b78b0af79bba1d1ea2156492e1b83dea8b8a6ae27395/propcache-0.4.1-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c07fda85708bc48578467e85099645167a955ba093be0a2dcba962195676e859", size = 273856, upload-time = "2025-10-08T19:47:34.906Z" }, + { url = "https://files.pythonhosted.org/packages/36/1d/fc272a63c8d3bbad6878c336c7a7dea15e8f2d23a544bda43205dfa83ada/propcache-0.4.1-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:af223b406d6d000830c6f65f1e6431783fc3f713ba3e6cc8c024d5ee96170a4b", size = 280420, upload-time = "2025-10-08T19:47:36.338Z" }, + { url = "https://files.pythonhosted.org/packages/07/0c/01f2219d39f7e53d52e5173bcb09c976609ba30209912a0680adfb8c593a/propcache-0.4.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a78372c932c90ee474559c5ddfffd718238e8673c340dc21fe45c5b8b54559a0", size = 263254, upload-time = "2025-10-08T19:47:37.692Z" }, + { url = "https://files.pythonhosted.org/packages/2d/18/cd28081658ce597898f0c4d174d4d0f3c5b6d4dc27ffafeef835c95eb359/propcache-0.4.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:564d9f0d4d9509e1a870c920a89b2fec951b44bf5ba7d537a9e7c1ccec2c18af", size = 261205, upload-time = "2025-10-08T19:47:39.659Z" }, + { url = "https://files.pythonhosted.org/packages/7a/71/1f9e22eb8b8316701c2a19fa1f388c8a3185082607da8e406a803c9b954e/propcache-0.4.1-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:17612831fda0138059cc5546f4d12a2aacfb9e47068c06af35c400ba58ba7393", size = 247873, upload-time = "2025-10-08T19:47:41.084Z" }, + { url = "https://files.pythonhosted.org/packages/4a/65/3d4b61f36af2b4eddba9def857959f1016a51066b4f1ce348e0cf7881f58/propcache-0.4.1-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:41a89040cb10bd345b3c1a873b2bf36413d48da1def52f268a055f7398514874", size = 262739, upload-time = "2025-10-08T19:47:42.51Z" }, + { url = "https://files.pythonhosted.org/packages/2a/42/26746ab087faa77c1c68079b228810436ccd9a5ce9ac85e2b7307195fd06/propcache-0.4.1-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:e35b88984e7fa64aacecea39236cee32dd9bd8c55f57ba8a75cf2399553f9bd7", size = 263514, upload-time = "2025-10-08T19:47:43.927Z" }, + { url = "https://files.pythonhosted.org/packages/94/13/630690fe201f5502d2403dd3cfd451ed8858fe3c738ee88d095ad2ff407b/propcache-0.4.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6f8b465489f927b0df505cbe26ffbeed4d6d8a2bbc61ce90eb074ff129ef0ab1", size = 257781, upload-time = "2025-10-08T19:47:45.448Z" }, + { url = "https://files.pythonhosted.org/packages/92/f7/1d4ec5841505f423469efbfc381d64b7b467438cd5a4bbcbb063f3b73d27/propcache-0.4.1-cp313-cp313t-win32.whl", hash = "sha256:2ad890caa1d928c7c2965b48f3a3815c853180831d0e5503d35cf00c472f4717", size = 41396, upload-time = "2025-10-08T19:47:47.202Z" }, + { url = "https://files.pythonhosted.org/packages/48/f0/615c30622316496d2cbbc29f5985f7777d3ada70f23370608c1d3e081c1f/propcache-0.4.1-cp313-cp313t-win_amd64.whl", hash = "sha256:f7ee0e597f495cf415bcbd3da3caa3bd7e816b74d0d52b8145954c5e6fd3ff37", size = 44897, upload-time = "2025-10-08T19:47:48.336Z" }, + { url = "https://files.pythonhosted.org/packages/fd/ca/6002e46eccbe0e33dcd4069ef32f7f1c9e243736e07adca37ae8c4830ec3/propcache-0.4.1-cp313-cp313t-win_arm64.whl", hash = "sha256:929d7cbe1f01bb7baffb33dc14eb5691c95831450a26354cd210a8155170c93a", size = 39789, upload-time = "2025-10-08T19:47:49.876Z" }, + { url = "https://files.pythonhosted.org/packages/8e/5c/bca52d654a896f831b8256683457ceddd490ec18d9ec50e97dfd8fc726a8/propcache-0.4.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:3f7124c9d820ba5548d431afb4632301acf965db49e666aa21c305cbe8c6de12", size = 78152, upload-time = "2025-10-08T19:47:51.051Z" }, + { url = "https://files.pythonhosted.org/packages/65/9b/03b04e7d82a5f54fb16113d839f5ea1ede58a61e90edf515f6577c66fa8f/propcache-0.4.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:c0d4b719b7da33599dfe3b22d3db1ef789210a0597bc650b7cee9c77c2be8c5c", size = 44869, upload-time = "2025-10-08T19:47:52.594Z" }, + { url = "https://files.pythonhosted.org/packages/b2/fa/89a8ef0468d5833a23fff277b143d0573897cf75bd56670a6d28126c7d68/propcache-0.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:9f302f4783709a78240ebc311b793f123328716a60911d667e0c036bc5dcbded", size = 46596, upload-time = "2025-10-08T19:47:54.073Z" }, + { url = "https://files.pythonhosted.org/packages/86/bd/47816020d337f4a746edc42fe8d53669965138f39ee117414c7d7a340cfe/propcache-0.4.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c80ee5802e3fb9ea37938e7eecc307fb984837091d5fd262bb37238b1ae97641", size = 206981, upload-time = "2025-10-08T19:47:55.715Z" }, + { url = "https://files.pythonhosted.org/packages/df/f6/c5fa1357cc9748510ee55f37173eb31bfde6d94e98ccd9e6f033f2fc06e1/propcache-0.4.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ed5a841e8bb29a55fb8159ed526b26adc5bdd7e8bd7bf793ce647cb08656cdf4", size = 211490, upload-time = "2025-10-08T19:47:57.499Z" }, + { url = "https://files.pythonhosted.org/packages/80/1e/e5889652a7c4a3846683401a48f0f2e5083ce0ec1a8a5221d8058fbd1adf/propcache-0.4.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:55c72fd6ea2da4c318e74ffdf93c4fe4e926051133657459131a95c846d16d44", size = 215371, upload-time = "2025-10-08T19:47:59.317Z" }, + { url = "https://files.pythonhosted.org/packages/b2/f2/889ad4b2408f72fe1a4f6a19491177b30ea7bf1a0fd5f17050ca08cfc882/propcache-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8326e144341460402713f91df60ade3c999d601e7eb5ff8f6f7862d54de0610d", size = 201424, upload-time = "2025-10-08T19:48:00.67Z" }, + { url = "https://files.pythonhosted.org/packages/27/73/033d63069b57b0812c8bd19f311faebeceb6ba31b8f32b73432d12a0b826/propcache-0.4.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:060b16ae65bc098da7f6d25bf359f1f31f688384858204fe5d652979e0015e5b", size = 197566, upload-time = "2025-10-08T19:48:02.604Z" }, + { url = "https://files.pythonhosted.org/packages/dc/89/ce24f3dc182630b4e07aa6d15f0ff4b14ed4b9955fae95a0b54c58d66c05/propcache-0.4.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:89eb3fa9524f7bec9de6e83cf3faed9d79bffa560672c118a96a171a6f55831e", size = 193130, upload-time = "2025-10-08T19:48:04.499Z" }, + { url = "https://files.pythonhosted.org/packages/a9/24/ef0d5fd1a811fb5c609278d0209c9f10c35f20581fcc16f818da959fc5b4/propcache-0.4.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:dee69d7015dc235f526fe80a9c90d65eb0039103fe565776250881731f06349f", size = 202625, upload-time = "2025-10-08T19:48:06.213Z" }, + { url = "https://files.pythonhosted.org/packages/f5/02/98ec20ff5546f68d673df2f7a69e8c0d076b5abd05ca882dc7ee3a83653d/propcache-0.4.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:5558992a00dfd54ccbc64a32726a3357ec93825a418a401f5cc67df0ac5d9e49", size = 204209, upload-time = "2025-10-08T19:48:08.432Z" }, + { url = "https://files.pythonhosted.org/packages/a0/87/492694f76759b15f0467a2a93ab68d32859672b646aa8a04ce4864e7932d/propcache-0.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:c9b822a577f560fbd9554812526831712c1436d2c046cedee4c3796d3543b144", size = 197797, upload-time = "2025-10-08T19:48:09.968Z" }, + { url = "https://files.pythonhosted.org/packages/ee/36/66367de3575db1d2d3f3d177432bd14ee577a39d3f5d1b3d5df8afe3b6e2/propcache-0.4.1-cp314-cp314-win32.whl", hash = "sha256:ab4c29b49d560fe48b696cdcb127dd36e0bc2472548f3bf56cc5cb3da2b2984f", size = 38140, upload-time = "2025-10-08T19:48:11.232Z" }, + { url = "https://files.pythonhosted.org/packages/0c/2a/a758b47de253636e1b8aef181c0b4f4f204bf0dd964914fb2af90a95b49b/propcache-0.4.1-cp314-cp314-win_amd64.whl", hash = "sha256:5a103c3eb905fcea0ab98be99c3a9a5ab2de60228aa5aceedc614c0281cf6153", size = 41257, upload-time = "2025-10-08T19:48:12.707Z" }, + { url = "https://files.pythonhosted.org/packages/34/5e/63bd5896c3fec12edcbd6f12508d4890d23c265df28c74b175e1ef9f4f3b/propcache-0.4.1-cp314-cp314-win_arm64.whl", hash = "sha256:74c1fb26515153e482e00177a1ad654721bf9207da8a494a0c05e797ad27b992", size = 38097, upload-time = "2025-10-08T19:48:13.923Z" }, + { url = "https://files.pythonhosted.org/packages/99/85/9ff785d787ccf9bbb3f3106f79884a130951436f58392000231b4c737c80/propcache-0.4.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:824e908bce90fb2743bd6b59db36eb4f45cd350a39637c9f73b1c1ea66f5b75f", size = 81455, upload-time = "2025-10-08T19:48:15.16Z" }, + { url = "https://files.pythonhosted.org/packages/90/85/2431c10c8e7ddb1445c1f7c4b54d886e8ad20e3c6307e7218f05922cad67/propcache-0.4.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:c2b5e7db5328427c57c8e8831abda175421b709672f6cfc3d630c3b7e2146393", size = 46372, upload-time = "2025-10-08T19:48:16.424Z" }, + { url = "https://files.pythonhosted.org/packages/01/20/b0972d902472da9bcb683fa595099911f4d2e86e5683bcc45de60dd05dc3/propcache-0.4.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6f6ff873ed40292cd4969ef5310179afd5db59fdf055897e282485043fc80ad0", size = 48411, upload-time = "2025-10-08T19:48:17.577Z" }, + { url = "https://files.pythonhosted.org/packages/e2/e3/7dc89f4f21e8f99bad3d5ddb3a3389afcf9da4ac69e3deb2dcdc96e74169/propcache-0.4.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:49a2dc67c154db2c1463013594c458881a069fcf98940e61a0569016a583020a", size = 275712, upload-time = "2025-10-08T19:48:18.901Z" }, + { url = "https://files.pythonhosted.org/packages/20/67/89800c8352489b21a8047c773067644e3897f02ecbbd610f4d46b7f08612/propcache-0.4.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:005f08e6a0529984491e37d8dbc3dd86f84bd78a8ceb5fa9a021f4c48d4984be", size = 273557, upload-time = "2025-10-08T19:48:20.762Z" }, + { url = "https://files.pythonhosted.org/packages/e2/a1/b52b055c766a54ce6d9c16d9aca0cad8059acd9637cdf8aa0222f4a026ef/propcache-0.4.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5c3310452e0d31390da9035c348633b43d7e7feb2e37be252be6da45abd1abcc", size = 280015, upload-time = "2025-10-08T19:48:22.592Z" }, + { url = "https://files.pythonhosted.org/packages/48/c8/33cee30bd890672c63743049f3c9e4be087e6780906bfc3ec58528be59c1/propcache-0.4.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c3c70630930447f9ef1caac7728c8ad1c56bc5015338b20fed0d08ea2480b3a", size = 262880, upload-time = "2025-10-08T19:48:23.947Z" }, + { url = "https://files.pythonhosted.org/packages/0c/b1/8f08a143b204b418285c88b83d00edbd61afbc2c6415ffafc8905da7038b/propcache-0.4.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8e57061305815dfc910a3634dcf584f08168a8836e6999983569f51a8544cd89", size = 260938, upload-time = "2025-10-08T19:48:25.656Z" }, + { url = "https://files.pythonhosted.org/packages/cf/12/96e4664c82ca2f31e1c8dff86afb867348979eb78d3cb8546a680287a1e9/propcache-0.4.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:521a463429ef54143092c11a77e04056dd00636f72e8c45b70aaa3140d639726", size = 247641, upload-time = "2025-10-08T19:48:27.207Z" }, + { url = "https://files.pythonhosted.org/packages/18/ed/e7a9cfca28133386ba52278136d42209d3125db08d0a6395f0cba0c0285c/propcache-0.4.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:120c964da3fdc75e3731aa392527136d4ad35868cc556fd09bb6d09172d9a367", size = 262510, upload-time = "2025-10-08T19:48:28.65Z" }, + { url = "https://files.pythonhosted.org/packages/f5/76/16d8bf65e8845dd62b4e2b57444ab81f07f40caa5652b8969b87ddcf2ef6/propcache-0.4.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:d8f353eb14ee3441ee844ade4277d560cdd68288838673273b978e3d6d2c8f36", size = 263161, upload-time = "2025-10-08T19:48:30.133Z" }, + { url = "https://files.pythonhosted.org/packages/e7/70/c99e9edb5d91d5ad8a49fa3c1e8285ba64f1476782fed10ab251ff413ba1/propcache-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ab2943be7c652f09638800905ee1bab2c544e537edb57d527997a24c13dc1455", size = 257393, upload-time = "2025-10-08T19:48:31.567Z" }, + { url = "https://files.pythonhosted.org/packages/08/02/87b25304249a35c0915d236575bc3574a323f60b47939a2262b77632a3ee/propcache-0.4.1-cp314-cp314t-win32.whl", hash = "sha256:05674a162469f31358c30bcaa8883cb7829fa3110bf9c0991fe27d7896c42d85", size = 42546, upload-time = "2025-10-08T19:48:32.872Z" }, + { url = "https://files.pythonhosted.org/packages/cb/ef/3c6ecf8b317aa982f309835e8f96987466123c6e596646d4e6a1dfcd080f/propcache-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:990f6b3e2a27d683cb7602ed6c86f15ee6b43b1194736f9baaeb93d0016633b1", size = 46259, upload-time = "2025-10-08T19:48:34.226Z" }, + { url = "https://files.pythonhosted.org/packages/c4/2d/346e946d4951f37eca1e4f55be0f0174c52cd70720f84029b02f296f4a38/propcache-0.4.1-cp314-cp314t-win_arm64.whl", hash = "sha256:ecef2343af4cc68e05131e45024ba34f6095821988a9d0a02aa7c73fcc448aa9", size = 40428, upload-time = "2025-10-08T19:48:35.441Z" }, + { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, +] + +[[package]] +name = "protobuf" +version = "6.33.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/25/7c72c307aafc96fa87062aa6291d9f7c94836e43214d43722e86037aac02/protobuf-6.33.5.tar.gz", hash = "sha256:6ddcac2a081f8b7b9642c09406bc6a4290128fce5f471cddd165960bb9119e5c", size = 444465, upload-time = "2026-01-29T21:51:33.494Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b1/79/af92d0a8369732b027e6d6084251dd8e782c685c72da161bd4a2e00fbabb/protobuf-6.33.5-cp310-abi3-win32.whl", hash = "sha256:d71b040839446bac0f4d162e758bea99c8251161dae9d0983a3b88dee345153b", size = 425769, upload-time = "2026-01-29T21:51:21.751Z" }, + { url = "https://files.pythonhosted.org/packages/55/75/bb9bc917d10e9ee13dee8607eb9ab963b7cf8be607c46e7862c748aa2af7/protobuf-6.33.5-cp310-abi3-win_amd64.whl", hash = "sha256:3093804752167bcab3998bec9f1048baae6e29505adaf1afd14a37bddede533c", size = 437118, upload-time = "2026-01-29T21:51:24.022Z" }, + { url = "https://files.pythonhosted.org/packages/a2/6b/e48dfc1191bc5b52950246275bf4089773e91cb5ba3592621723cdddca62/protobuf-6.33.5-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:a5cb85982d95d906df1e2210e58f8e4f1e3cdc088e52c921a041f9c9a0386de5", size = 427766, upload-time = "2026-01-29T21:51:25.413Z" }, + { url = "https://files.pythonhosted.org/packages/4e/b1/c79468184310de09d75095ed1314b839eb2f72df71097db9d1404a1b2717/protobuf-6.33.5-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:9b71e0281f36f179d00cbcb119cb19dec4d14a81393e5ea220f64b286173e190", size = 324638, upload-time = "2026-01-29T21:51:26.423Z" }, + { url = "https://files.pythonhosted.org/packages/c5/f5/65d838092fd01c44d16037953fd4c2cc851e783de9b8f02b27ec4ffd906f/protobuf-6.33.5-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:8afa18e1d6d20af15b417e728e9f60f3aa108ee76f23c3b2c07a2c3b546d3afd", size = 339411, upload-time = "2026-01-29T21:51:27.446Z" }, + { url = "https://files.pythonhosted.org/packages/9b/53/a9443aa3ca9ba8724fdfa02dd1887c1bcd8e89556b715cfbacca6b63dbec/protobuf-6.33.5-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:cbf16ba3350fb7b889fca858fb215967792dc125b35c7976ca4818bee3521cf0", size = 323465, upload-time = "2026-01-29T21:51:28.925Z" }, + { url = "https://files.pythonhosted.org/packages/57/bf/2086963c69bdac3d7cff1cc7ff79b8ce5ea0bec6797a017e1be338a46248/protobuf-6.33.5-py3-none-any.whl", hash = "sha256:69915a973dd0f60f31a08b8318b73eab2bd6a392c79184b3612226b0a3f8ec02", size = 170687, upload-time = "2026-01-29T21:51:32.557Z" }, +] + +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" }, + { url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" }, + { url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" }, + { url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" }, + { url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" }, + { url = "https://files.pythonhosted.org/packages/81/69/ef179ab5ca24f32acc1dac0c247fd6a13b501fd5534dbae0e05a1c48b66d/psutil-7.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00", size = 130664, upload-time = "2026-01-28T18:15:09.469Z" }, + { url = "https://files.pythonhosted.org/packages/7b/64/665248b557a236d3fa9efc378d60d95ef56dd0a490c2cd37dafc7660d4a9/psutil-7.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9", size = 131087, upload-time = "2026-01-28T18:15:11.724Z" }, + { url = "https://files.pythonhosted.org/packages/d5/2e/e6782744700d6759ebce3043dcfa661fb61e2fb752b91cdeae9af12c2178/psutil-7.2.2-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a", size = 182383, upload-time = "2026-01-28T18:15:13.445Z" }, + { url = "https://files.pythonhosted.org/packages/57/49/0a41cefd10cb7505cdc04dab3eacf24c0c2cb158a998b8c7b1d27ee2c1f5/psutil-7.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf", size = 185210, upload-time = "2026-01-28T18:15:16.002Z" }, + { url = "https://files.pythonhosted.org/packages/dd/2c/ff9bfb544f283ba5f83ba725a3c5fec6d6b10b8f27ac1dc641c473dc390d/psutil-7.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1", size = 141228, upload-time = "2026-01-28T18:15:18.385Z" }, + { url = "https://files.pythonhosted.org/packages/f2/fc/f8d9c31db14fcec13748d373e668bc3bed94d9077dbc17fb0eebc073233c/psutil-7.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841", size = 136284, upload-time = "2026-01-28T18:15:19.912Z" }, + { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, + { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, + { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, + { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, + { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, +] + +[[package]] +name = "pybind11" +version = "3.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2f/7b/a6d8dcb83c457e24a9df1e4d8fd5fb8034d4bbc62f3c324681e8a9ba57c2/pybind11-3.0.1.tar.gz", hash = "sha256:9c0f40056a016da59bab516efb523089139fcc6f2ba7e4930854c61efb932051", size = 546914, upload-time = "2025-08-22T20:09:27.265Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cd/8a/37362fc2b949d5f733a8b0f2ff51ba423914cabefe69f1d1b6aab710f5fe/pybind11-3.0.1-py3-none-any.whl", hash = "sha256:aa8f0aa6e0a94d3b64adfc38f560f33f15e589be2175e103c0a33c6bce55ee89", size = 293611, upload-time = "2025-08-22T20:09:25.235Z" }, +] + +[[package]] +name = "pygments" +version = "2.19.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b0/77/a5b8c569bf593b0140bde72ea885a803b82086995367bf2037de0159d924/pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887", size = 4968631, upload-time = "2025-06-21T13:39:12.283Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" }, +] + +[[package]] +name = "pyparsing" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/91/9c6ee907786a473bf81c5f53cf703ba0957b23ab84c264080fb5a450416f/pyparsing-3.3.2.tar.gz", hash = "sha256:c777f4d763f140633dcb6d8a3eda953bf7a214dc4eff598413c070bcdc117cbc", size = 6851574, upload-time = "2026-01-21T03:57:59.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, +] + +[[package]] +name = "pysocks" +version = "1.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bd/11/293dd436aea955d45fc4e8a35b6ae7270f5b8e00b53cf6c024c83b657a11/PySocks-1.7.1.tar.gz", hash = "sha256:3f8804571ebe159c380ac6de37643bb4685970655d3bba243530d6558b799aa0", size = 284429, upload-time = "2019-09-20T02:07:35.714Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/59/b4572118e098ac8e46e399a1dd0f2d85403ce8bbaad9ec79373ed6badaf9/PySocks-1.7.1-py3-none-any.whl", hash = "sha256:2725bd0a9925919b9b51739eea5f9e2bae91e83288108a9ad338b2e3a4435ee5", size = 16725, upload-time = "2019-09-20T02:06:22.938Z" }, +] + +[[package]] +name = "pytest" +version = "9.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" }, +] + +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + +[[package]] +name = "python-dotenv" +version = "1.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f0/26/19cadc79a718c5edbec86fd4919a6b6d3f681039a2f6d66d14be94e75fb9/python_dotenv-1.2.1.tar.gz", hash = "sha256:42667e897e16ab0d66954af0e60a9caa94f0fd4ecf3aaf6d2d260eec1aa36ad6", size = 44221, upload-time = "2025-10-26T15:12:10.434Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/1b/a298b06749107c305e1fe0f814c6c74aea7b2f1e10989cb30f544a1b3253/python_dotenv-1.2.1-py3-none-any.whl", hash = "sha256:b81ee9561e9ca4004139c6cbba3a238c32b03e4894671e181b671e8cb8425d61", size = 21230, upload-time = "2025-10-26T15:12:09.109Z" }, +] + +[[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/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8", size = 181669, upload-time = "2025-09-25T21:32:23.673Z" }, + { url = "https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1", size = 173252, upload-time = "2025-09-25T21:32:25.149Z" }, + { url = "https://files.pythonhosted.org/packages/50/31/b20f376d3f810b9b2371e72ef5adb33879b25edb7a6d072cb7ca0c486398/pyyaml-6.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c", size = 767081, upload-time = "2025-09-25T21:32:26.575Z" }, + { url = "https://files.pythonhosted.org/packages/49/1e/a55ca81e949270d5d4432fbbd19dfea5321eda7c41a849d443dc92fd1ff7/pyyaml-6.0.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5", size = 841159, upload-time = "2025-09-25T21:32:27.727Z" }, + { url = "https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6", size = 801626, upload-time = "2025-09-25T21:32:28.878Z" }, + { url = "https://files.pythonhosted.org/packages/f9/11/ba845c23988798f40e52ba45f34849aa8a1f2d4af4b798588010792ebad6/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6", size = 753613, upload-time = "2025-09-25T21:32:30.178Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e0/7966e1a7bfc0a45bf0a7fb6b98ea03fc9b8d84fa7f2229e9659680b69ee3/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be", size = 794115, upload-time = "2025-09-25T21:32:31.353Z" }, + { url = "https://files.pythonhosted.org/packages/de/94/980b50a6531b3019e45ddeada0626d45fa85cbe22300844a7983285bed3b/pyyaml-6.0.3-cp313-cp313-win32.whl", hash = "sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26", size = 137427, upload-time = "2025-09-25T21:32:32.58Z" }, + { url = "https://files.pythonhosted.org/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c", size = 154090, upload-time = "2025-09-25T21:32:33.659Z" }, + { url = "https://files.pythonhosted.org/packages/73/e8/2bdf3ca2090f68bb3d75b44da7bbc71843b19c9f2b9cb9b0f4ab7a5a4329/pyyaml-6.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb", size = 140246, upload-time = "2025-09-25T21:32:34.663Z" }, + { url = "https://files.pythonhosted.org/packages/9d/8c/f4bd7f6465179953d3ac9bc44ac1a8a3e6122cf8ada906b4f96c60172d43/pyyaml-6.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:8d1fab6bb153a416f9aeb4b8763bc0f22a5586065f86f7664fc23339fc1c1fac", size = 181814, upload-time = "2025-09-25T21:32:35.712Z" }, + { url = "https://files.pythonhosted.org/packages/bd/9c/4d95bb87eb2063d20db7b60faa3840c1b18025517ae857371c4dd55a6b3a/pyyaml-6.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:34d5fcd24b8445fadc33f9cf348c1047101756fd760b4dacb5c3e99755703310", size = 173809, upload-time = "2025-09-25T21:32:36.789Z" }, + { url = "https://files.pythonhosted.org/packages/92/b5/47e807c2623074914e29dabd16cbbdd4bf5e9b2db9f8090fa64411fc5382/pyyaml-6.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7", size = 766454, upload-time = "2025-09-25T21:32:37.966Z" }, + { url = "https://files.pythonhosted.org/packages/02/9e/e5e9b168be58564121efb3de6859c452fccde0ab093d8438905899a3a483/pyyaml-6.0.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b3bc83488de33889877a0f2543ade9f70c67d66d9ebb4ac959502e12de895788", size = 836355, upload-time = "2025-09-25T21:32:39.178Z" }, + { url = "https://files.pythonhosted.org/packages/88/f9/16491d7ed2a919954993e48aa941b200f38040928474c9e85ea9e64222c3/pyyaml-6.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5", size = 794175, upload-time = "2025-09-25T21:32:40.865Z" }, + { url = "https://files.pythonhosted.org/packages/dd/3f/5989debef34dc6397317802b527dbbafb2b4760878a53d4166579111411e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764", size = 755228, upload-time = "2025-09-25T21:32:42.084Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ce/af88a49043cd2e265be63d083fc75b27b6ed062f5f9fd6cdc223ad62f03e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35", size = 789194, upload-time = "2025-09-25T21:32:43.362Z" }, + { url = "https://files.pythonhosted.org/packages/23/20/bb6982b26a40bb43951265ba29d4c246ef0ff59c9fdcdf0ed04e0687de4d/pyyaml-6.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:4a2e8cebe2ff6ab7d1050ecd59c25d4c8bd7e6f400f5f82b96557ac0abafd0ac", size = 156429, upload-time = "2025-09-25T21:32:57.844Z" }, + { url = "https://files.pythonhosted.org/packages/f4/f4/a4541072bb9422c8a883ab55255f918fa378ecf083f5b85e87fc2b4eda1b/pyyaml-6.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:93dda82c9c22deb0a405ea4dc5f2d0cda384168e466364dec6255b293923b2f3", size = 143912, upload-time = "2025-09-25T21:32:59.247Z" }, + { url = "https://files.pythonhosted.org/packages/7c/f9/07dd09ae774e4616edf6cda684ee78f97777bdd15847253637a6f052a62f/pyyaml-6.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:02893d100e99e03eda1c8fd5c441d8c60103fd175728e23e431db1b589cf5ab3", size = 189108, upload-time = "2025-09-25T21:32:44.377Z" }, + { url = "https://files.pythonhosted.org/packages/4e/78/8d08c9fb7ce09ad8c38ad533c1191cf27f7ae1effe5bb9400a46d9437fcf/pyyaml-6.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c1ff362665ae507275af2853520967820d9124984e0f7466736aea23d8611fba", size = 183641, upload-time = "2025-09-25T21:32:45.407Z" }, + { url = "https://files.pythonhosted.org/packages/7b/5b/3babb19104a46945cf816d047db2788bcaf8c94527a805610b0289a01c6b/pyyaml-6.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c", size = 831901, upload-time = "2025-09-25T21:32:48.83Z" }, + { url = "https://files.pythonhosted.org/packages/8b/cc/dff0684d8dc44da4d22a13f35f073d558c268780ce3c6ba1b87055bb0b87/pyyaml-6.0.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a80cb027f6b349846a3bf6d73b5e95e782175e52f22108cfa17876aaeff93702", size = 861132, upload-time = "2025-09-25T21:32:50.149Z" }, + { url = "https://files.pythonhosted.org/packages/b1/5e/f77dc6b9036943e285ba76b49e118d9ea929885becb0a29ba8a7c75e29fe/pyyaml-6.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c", size = 839261, upload-time = "2025-09-25T21:32:51.808Z" }, + { url = "https://files.pythonhosted.org/packages/ce/88/a9db1376aa2a228197c58b37302f284b5617f56a5d959fd1763fb1675ce6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065", size = 805272, upload-time = "2025-09-25T21:32:52.941Z" }, + { url = "https://files.pythonhosted.org/packages/da/92/1446574745d74df0c92e6aa4a7b0b3130706a4142b2d1a5869f2eaa423c6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65", size = 829923, upload-time = "2025-09-25T21:32:54.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/7a/1c7270340330e575b92f397352af856a8c06f230aa3e76f86b39d01b416a/pyyaml-6.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4ad1906908f2f5ae4e5a8ddfce73c320c2a1429ec52eafd27138b7f1cbe341c9", size = 174062, upload-time = "2025-09-25T21:32:55.767Z" }, + { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, +] + +[[package]] +name = "requests" +version = "2.32.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c9/74/b3ff8e6c8446842c3f5c837e9c3dfcfe2018ea6ecef224c710c85ef728f4/requests-2.32.5.tar.gz", hash = "sha256:dbba0bac56e100853db0ea71b82b4dfd5fe2bf6d3754a8893c3af500cec7d7cf", size = 134517, upload-time = "2025-08-18T20:46:02.573Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" }, +] + +[package.optional-dependencies] +socks = [ + { name = "pysocks" }, +] + +[[package]] +name = "rich" +version = "14.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/74/99/a4cab2acbb884f80e558b0771e97e21e939c5dfb460f488d19df485e8298/rich-14.3.2.tar.gz", hash = "sha256:e712f11c1a562a11843306f5ed999475f09ac31ffb64281f73ab29ffdda8b3b8", size = 230143, upload-time = "2026-02-01T16:20:47.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ef/45/615f5babd880b4bd7d405cc0dc348234c5ffb6ed1ea33e152ede08b2072d/rich-14.3.2-py3-none-any.whl", hash = "sha256:08e67c3e90884651da3239ea668222d19bea7b589149d8014a21c633420dbb69", size = 309963, upload-time = "2026-02-01T16:20:46.078Z" }, +] + +[[package]] +name = "ruff" +version = "0.15.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/da/31/d6e536cdebb6568ae75a7f00e4b4819ae0ad2640c3604c305a0428680b0c/ruff-0.15.4.tar.gz", hash = "sha256:3412195319e42d634470cc97aa9803d07e9d5c9223b99bcb1518f0c725f26ae1", size = 4569550, upload-time = "2026-02-26T20:04:14.959Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/82/c11a03cfec3a4d26a0ea1e571f0f44be5993b923f905eeddfc397c13d360/ruff-0.15.4-py3-none-linux_armv6l.whl", hash = "sha256:a1810931c41606c686bae8b5b9a8072adac2f611bb433c0ba476acba17a332e0", size = 10453333, upload-time = "2026-02-26T20:04:20.093Z" }, + { url = "https://files.pythonhosted.org/packages/ce/5d/6a1f271f6e31dffb31855996493641edc3eef8077b883eaf007a2f1c2976/ruff-0.15.4-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:5a1632c66672b8b4d3e1d1782859e98d6e0b4e70829530666644286600a33992", size = 10853356, upload-time = "2026-02-26T20:04:05.808Z" }, + { url = "https://files.pythonhosted.org/packages/b1/d8/0fab9f8842b83b1a9c2bf81b85063f65e93fb512e60effa95b0be49bfc54/ruff-0.15.4-py3-none-macosx_11_0_arm64.whl", hash = "sha256:a4386ba2cd6c0f4ff75252845906acc7c7c8e1ac567b7bc3d373686ac8c222ba", size = 10187434, upload-time = "2026-02-26T20:03:54.656Z" }, + { url = "https://files.pythonhosted.org/packages/85/cc/cc220fd9394eff5db8d94dec199eec56dd6c9f3651d8869d024867a91030/ruff-0.15.4-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b2496488bdfd3732747558b6f95ae427ff066d1fcd054daf75f5a50674411e75", size = 10535456, upload-time = "2026-02-26T20:03:52.738Z" }, + { url = "https://files.pythonhosted.org/packages/fa/0f/bced38fa5cf24373ec767713c8e4cadc90247f3863605fb030e597878661/ruff-0.15.4-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3f1c4893841ff2d54cbda1b2860fa3260173df5ddd7b95d370186f8a5e66a4ac", size = 10287772, upload-time = "2026-02-26T20:04:08.138Z" }, + { url = "https://files.pythonhosted.org/packages/2b/90/58a1802d84fed15f8f281925b21ab3cecd813bde52a8ca033a4de8ab0e7a/ruff-0.15.4-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:820b8766bd65503b6c30aaa6331e8ef3a6e564f7999c844e9a547c40179e440a", size = 11049051, upload-time = "2026-02-26T20:04:03.53Z" }, + { url = "https://files.pythonhosted.org/packages/d2/ac/b7ad36703c35f3866584564dc15f12f91cb1a26a897dc2fd13d7cb3ae1af/ruff-0.15.4-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c9fb74bab47139c1751f900f857fa503987253c3ef89129b24ed375e72873e85", size = 11890494, upload-time = "2026-02-26T20:04:10.497Z" }, + { url = "https://files.pythonhosted.org/packages/93/3d/3eb2f47a39a8b0da99faf9c54d3eb24720add1e886a5309d4d1be73a6380/ruff-0.15.4-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f80c98765949c518142b3a50a5db89343aa90f2c2bf7799de9986498ae6176db", size = 11326221, upload-time = "2026-02-26T20:04:12.84Z" }, + { url = "https://files.pythonhosted.org/packages/ff/90/bf134f4c1e5243e62690e09d63c55df948a74084c8ac3e48a88468314da6/ruff-0.15.4-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:451a2e224151729b3b6c9ffb36aed9091b2996fe4bdbd11f47e27d8f2e8888ec", size = 11168459, upload-time = "2026-02-26T20:04:00.969Z" }, + { url = "https://files.pythonhosted.org/packages/b5/e5/a64d27688789b06b5d55162aafc32059bb8c989c61a5139a36e1368285eb/ruff-0.15.4-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:a8f157f2e583c513c4f5f896163a93198297371f34c04220daf40d133fdd4f7f", size = 11104366, upload-time = "2026-02-26T20:03:48.099Z" }, + { url = "https://files.pythonhosted.org/packages/f1/f6/32d1dcb66a2559763fc3027bdd65836cad9eb09d90f2ed6a63d8e9252b02/ruff-0.15.4-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:917cc68503357021f541e69b35361c99387cdbbf99bd0ea4aa6f28ca99ff5338", size = 10510887, upload-time = "2026-02-26T20:03:45.771Z" }, + { url = "https://files.pythonhosted.org/packages/ff/92/22d1ced50971c5b6433aed166fcef8c9343f567a94cf2b9d9089f6aa80fe/ruff-0.15.4-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:e9737c8161da79fd7cfec19f1e35620375bd8b2a50c3e77fa3d2c16f574105cc", size = 10285939, upload-time = "2026-02-26T20:04:22.42Z" }, + { url = "https://files.pythonhosted.org/packages/e6/f4/7c20aec3143837641a02509a4668fb146a642fd1211846634edc17eb5563/ruff-0.15.4-py3-none-musllinux_1_2_i686.whl", hash = "sha256:291258c917539e18f6ba40482fe31d6f5ac023994ee11d7bdafd716f2aab8a68", size = 10765471, upload-time = "2026-02-26T20:03:58.924Z" }, + { url = "https://files.pythonhosted.org/packages/d0/09/6d2f7586f09a16120aebdff8f64d962d7c4348313c77ebb29c566cefc357/ruff-0.15.4-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:3f83c45911da6f2cd5936c436cf86b9f09f09165f033a99dcf7477e34041cbc3", size = 11263382, upload-time = "2026-02-26T20:04:24.424Z" }, + { url = "https://files.pythonhosted.org/packages/1b/fa/2ef715a1cd329ef47c1a050e10dee91a9054b7ce2fcfdd6a06d139afb7ec/ruff-0.15.4-py3-none-win32.whl", hash = "sha256:65594a2d557d4ee9f02834fcdf0a28daa8b3b9f6cb2cb93846025a36db47ef22", size = 10506664, upload-time = "2026-02-26T20:03:50.56Z" }, + { url = "https://files.pythonhosted.org/packages/d0/a8/c688ef7e29983976820d18710f955751d9f4d4eb69df658af3d006e2ba3e/ruff-0.15.4-py3-none-win_amd64.whl", hash = "sha256:04196ad44f0df220c2ece5b0e959c2f37c777375ec744397d21d15b50a75264f", size = 11651048, upload-time = "2026-02-26T20:04:17.191Z" }, + { url = "https://files.pythonhosted.org/packages/3e/0a/9e1be9035b37448ce2e68c978f0591da94389ade5a5abafa4cf99985d1b2/ruff-0.15.4-py3-none-win_arm64.whl", hash = "sha256:60d5177e8cfc70e51b9c5fad936c634872a74209f934c1e79107d11787ad5453", size = 10966776, upload-time = "2026-02-26T20:03:56.908Z" }, +] + +[[package]] +name = "scipy" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/56/3e/9cca699f3486ce6bc12ff46dc2031f1ec8eb9ccc9a320fdaf925f1417426/scipy-1.17.0.tar.gz", hash = "sha256:2591060c8e648d8b96439e111ac41fd8342fdeff1876be2e19dea3fe8930454e", size = 30396830, upload-time = "2026-01-10T21:34:23.009Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/51/3468fdfd49387ddefee1636f5cf6d03ce603b75205bf439bbf0e62069bfd/scipy-1.17.0-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:65ec32f3d32dfc48c72df4291345dae4f048749bc8d5203ee0a3f347f96c5ce6", size = 31344101, upload-time = "2026-01-10T21:26:30.25Z" }, + { url = "https://files.pythonhosted.org/packages/b2/9a/9406aec58268d437636069419e6977af953d1e246df941d42d3720b7277b/scipy-1.17.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:1f9586a58039d7229ce77b52f8472c972448cded5736eaf102d5658bbac4c269", size = 27950385, upload-time = "2026-01-10T21:26:36.801Z" }, + { url = "https://files.pythonhosted.org/packages/4f/98/e7342709e17afdfd1b26b56ae499ef4939b45a23a00e471dfb5375eea205/scipy-1.17.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:9fad7d3578c877d606b1150135c2639e9de9cecd3705caa37b66862977cc3e72", size = 20122115, upload-time = "2026-01-10T21:26:42.107Z" }, + { url = "https://files.pythonhosted.org/packages/fd/0e/9eeeb5357a64fd157cbe0302c213517c541cc16b8486d82de251f3c68ede/scipy-1.17.0-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:423ca1f6584fc03936972b5f7c06961670dbba9f234e71676a7c7ccf938a0d61", size = 22442402, upload-time = "2026-01-10T21:26:48.029Z" }, + { url = "https://files.pythonhosted.org/packages/c9/10/be13397a0e434f98e0c79552b2b584ae5bb1c8b2be95db421533bbca5369/scipy-1.17.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe508b5690e9eaaa9467fc047f833af58f1152ae51a0d0aed67aa5801f4dd7d6", size = 32696338, upload-time = "2026-01-10T21:26:55.521Z" }, + { url = "https://files.pythonhosted.org/packages/63/1e/12fbf2a3bb240161651c94bb5cdd0eae5d4e8cc6eaeceb74ab07b12a753d/scipy-1.17.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6680f2dfd4f6182e7d6db161344537da644d1cf85cf293f015c60a17ecf08752", size = 34977201, upload-time = "2026-01-10T21:27:03.501Z" }, + { url = "https://files.pythonhosted.org/packages/19/5b/1a63923e23ccd20bd32156d7dd708af5bbde410daa993aa2500c847ab2d2/scipy-1.17.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:eec3842ec9ac9de5917899b277428886042a93db0b227ebbe3a333b64ec7643d", size = 34777384, upload-time = "2026-01-10T21:27:11.423Z" }, + { url = "https://files.pythonhosted.org/packages/39/22/b5da95d74edcf81e540e467202a988c50fef41bd2011f46e05f72ba07df6/scipy-1.17.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:d7425fcafbc09a03731e1bc05581f5fad988e48c6a861f441b7ab729a49a55ea", size = 37379586, upload-time = "2026-01-10T21:27:20.171Z" }, + { url = "https://files.pythonhosted.org/packages/b9/b6/8ac583d6da79e7b9e520579f03007cb006f063642afd6b2eeb16b890bf93/scipy-1.17.0-cp313-cp313-win_amd64.whl", hash = "sha256:87b411e42b425b84777718cc41516b8a7e0795abfa8e8e1d573bf0ef014f0812", size = 36287211, upload-time = "2026-01-10T21:28:43.122Z" }, + { url = "https://files.pythonhosted.org/packages/55/fb/7db19e0b3e52f882b420417644ec81dd57eeef1bd1705b6f689d8ff93541/scipy-1.17.0-cp313-cp313-win_arm64.whl", hash = "sha256:357ca001c6e37601066092e7c89cca2f1ce74e2a520ca78d063a6d2201101df2", size = 24312646, upload-time = "2026-01-10T21:28:49.893Z" }, + { url = "https://files.pythonhosted.org/packages/20/b6/7feaa252c21cc7aff335c6c55e1b90ab3e3306da3f048109b8b639b94648/scipy-1.17.0-cp313-cp313t-macosx_10_14_x86_64.whl", hash = "sha256:ec0827aa4d36cb79ff1b81de898e948a51ac0b9b1c43e4a372c0508c38c0f9a3", size = 31693194, upload-time = "2026-01-10T21:27:27.454Z" }, + { url = "https://files.pythonhosted.org/packages/76/bb/bbb392005abce039fb7e672cb78ac7d158700e826b0515cab6b5b60c26fb/scipy-1.17.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:819fc26862b4b3c73a60d486dbb919202f3d6d98c87cf20c223511429f2d1a97", size = 28365415, upload-time = "2026-01-10T21:27:34.26Z" }, + { url = "https://files.pythonhosted.org/packages/37/da/9d33196ecc99fba16a409c691ed464a3a283ac454a34a13a3a57c0d66f3a/scipy-1.17.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:363ad4ae2853d88ebcde3ae6ec46ccca903ea9835ee8ba543f12f575e7b07e4e", size = 20537232, upload-time = "2026-01-10T21:27:40.306Z" }, + { url = "https://files.pythonhosted.org/packages/56/9d/f4b184f6ddb28e9a5caea36a6f98e8ecd2a524f9127354087ce780885d83/scipy-1.17.0-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:979c3a0ff8e5ba254d45d59ebd38cde48fce4f10b5125c680c7a4bfe177aab07", size = 22791051, upload-time = "2026-01-10T21:27:46.539Z" }, + { url = "https://files.pythonhosted.org/packages/9b/9d/025cccdd738a72140efc582b1641d0dd4caf2e86c3fb127568dc80444e6e/scipy-1.17.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:130d12926ae34399d157de777472bf82e9061c60cc081372b3118edacafe1d00", size = 32815098, upload-time = "2026-01-10T21:27:54.389Z" }, + { url = "https://files.pythonhosted.org/packages/48/5f/09b879619f8bca15ce392bfc1894bd9c54377e01d1b3f2f3b595a1b4d945/scipy-1.17.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6e886000eb4919eae3a44f035e63f0fd8b651234117e8f6f29bad1cd26e7bc45", size = 35031342, upload-time = "2026-01-10T21:28:03.012Z" }, + { url = "https://files.pythonhosted.org/packages/f2/9a/f0f0a9f0aa079d2f106555b984ff0fbb11a837df280f04f71f056ea9c6e4/scipy-1.17.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:13c4096ac6bc31d706018f06a49abe0485f96499deb82066b94d19b02f664209", size = 34893199, upload-time = "2026-01-10T21:28:10.832Z" }, + { url = "https://files.pythonhosted.org/packages/90/b8/4f0f5cf0c5ea4d7548424e6533e6b17d164f34a6e2fb2e43ffebb6697b06/scipy-1.17.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:cacbaddd91fcffde703934897c5cd2c7cb0371fac195d383f4e1f1c5d3f3bd04", size = 37438061, upload-time = "2026-01-10T21:28:19.684Z" }, + { url = "https://files.pythonhosted.org/packages/f9/cc/2bd59140ed3b2fa2882fb15da0a9cb1b5a6443d67cfd0d98d4cec83a57ec/scipy-1.17.0-cp313-cp313t-win_amd64.whl", hash = "sha256:edce1a1cf66298cccdc48a1bdf8fb10a3bf58e8b58d6c3883dd1530e103f87c0", size = 36328593, upload-time = "2026-01-10T21:28:28.007Z" }, + { url = "https://files.pythonhosted.org/packages/13/1b/c87cc44a0d2c7aaf0f003aef2904c3d097b422a96c7e7c07f5efd9073c1b/scipy-1.17.0-cp313-cp313t-win_arm64.whl", hash = "sha256:30509da9dbec1c2ed8f168b8d8aa853bc6723fede1dbc23c7d43a56f5ab72a67", size = 24625083, upload-time = "2026-01-10T21:28:35.188Z" }, + { url = "https://files.pythonhosted.org/packages/1a/2d/51006cd369b8e7879e1c630999a19d1fbf6f8b5ed3e33374f29dc87e53b3/scipy-1.17.0-cp314-cp314-macosx_10_14_x86_64.whl", hash = "sha256:c17514d11b78be8f7e6331b983a65a7f5ca1fd037b95e27b280921fe5606286a", size = 31346803, upload-time = "2026-01-10T21:28:57.24Z" }, + { url = "https://files.pythonhosted.org/packages/d6/2e/2349458c3ce445f53a6c93d4386b1c4c5c0c540917304c01222ff95ff317/scipy-1.17.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:4e00562e519c09da34c31685f6acc3aa384d4d50604db0f245c14e1b4488bfa2", size = 27967182, upload-time = "2026-01-10T21:29:04.107Z" }, + { url = "https://files.pythonhosted.org/packages/5e/7c/df525fbfa77b878d1cfe625249529514dc02f4fd5f45f0f6295676a76528/scipy-1.17.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:f7df7941d71314e60a481e02d5ebcb3f0185b8d799c70d03d8258f6c80f3d467", size = 20139125, upload-time = "2026-01-10T21:29:10.179Z" }, + { url = "https://files.pythonhosted.org/packages/33/11/fcf9d43a7ed1234d31765ec643b0515a85a30b58eddccc5d5a4d12b5f194/scipy-1.17.0-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:aabf057c632798832f071a8dde013c2e26284043934f53b00489f1773b33527e", size = 22443554, upload-time = "2026-01-10T21:29:15.888Z" }, + { url = "https://files.pythonhosted.org/packages/80/5c/ea5d239cda2dd3d31399424967a24d556cf409fbea7b5b21412b0fd0a44f/scipy-1.17.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a38c3337e00be6fd8a95b4ed66b5d988bac4ec888fd922c2ea9fe5fb1603dd67", size = 32757834, upload-time = "2026-01-10T21:29:23.406Z" }, + { url = "https://files.pythonhosted.org/packages/b8/7e/8c917cc573310e5dc91cbeead76f1b600d3fb17cf0969db02c9cf92e3cfa/scipy-1.17.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00fb5f8ec8398ad90215008d8b6009c9db9fa924fd4c7d6be307c6f945f9cd73", size = 34995775, upload-time = "2026-01-10T21:29:31.915Z" }, + { url = "https://files.pythonhosted.org/packages/c5/43/176c0c3c07b3f7df324e7cdd933d3e2c4898ca202b090bd5ba122f9fe270/scipy-1.17.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:f2a4942b0f5f7c23c7cd641a0ca1955e2ae83dedcff537e3a0259096635e186b", size = 34841240, upload-time = "2026-01-10T21:29:39.995Z" }, + { url = "https://files.pythonhosted.org/packages/44/8c/d1f5f4b491160592e7f084d997de53a8e896a3ac01cd07e59f43ca222744/scipy-1.17.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:dbf133ced83889583156566d2bdf7a07ff89228fe0c0cb727f777de92092ec6b", size = 37394463, upload-time = "2026-01-10T21:29:48.723Z" }, + { url = "https://files.pythonhosted.org/packages/9f/ec/42a6657f8d2d087e750e9a5dde0b481fd135657f09eaf1cf5688bb23c338/scipy-1.17.0-cp314-cp314-win_amd64.whl", hash = "sha256:3625c631a7acd7cfd929e4e31d2582cf00f42fcf06011f59281271746d77e061", size = 37053015, upload-time = "2026-01-10T21:30:51.418Z" }, + { url = "https://files.pythonhosted.org/packages/27/58/6b89a6afd132787d89a362d443a7bddd511b8f41336a1ae47f9e4f000dc4/scipy-1.17.0-cp314-cp314-win_arm64.whl", hash = "sha256:9244608d27eafe02b20558523ba57f15c689357c85bdcfe920b1828750aa26eb", size = 24951312, upload-time = "2026-01-10T21:30:56.771Z" }, + { url = "https://files.pythonhosted.org/packages/e9/01/f58916b9d9ae0112b86d7c3b10b9e685625ce6e8248df139d0fcb17f7397/scipy-1.17.0-cp314-cp314t-macosx_10_14_x86_64.whl", hash = "sha256:2b531f57e09c946f56ad0b4a3b2abee778789097871fc541e267d2eca081cff1", size = 31706502, upload-time = "2026-01-10T21:29:56.326Z" }, + { url = "https://files.pythonhosted.org/packages/59/8e/2912a87f94a7d1f8b38aabc0faf74b82d3b6c9e22be991c49979f0eceed8/scipy-1.17.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:13e861634a2c480bd237deb69333ac79ea1941b94568d4b0efa5db5e263d4fd1", size = 28380854, upload-time = "2026-01-10T21:30:01.554Z" }, + { url = "https://files.pythonhosted.org/packages/bd/1c/874137a52dddab7d5d595c1887089a2125d27d0601fce8c0026a24a92a0b/scipy-1.17.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:eb2651271135154aa24f6481cbae5cc8af1f0dd46e6533fb7b56aa9727b6a232", size = 20552752, upload-time = "2026-01-10T21:30:05.93Z" }, + { url = "https://files.pythonhosted.org/packages/3f/f0/7518d171cb735f6400f4576cf70f756d5b419a07fe1867da34e2c2c9c11b/scipy-1.17.0-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:c5e8647f60679790c2f5c76be17e2e9247dc6b98ad0d3b065861e082c56e078d", size = 22803972, upload-time = "2026-01-10T21:30:10.651Z" }, + { url = "https://files.pythonhosted.org/packages/7c/74/3498563a2c619e8a3ebb4d75457486c249b19b5b04a30600dfd9af06bea5/scipy-1.17.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5fb10d17e649e1446410895639f3385fd2bf4c3c7dfc9bea937bddcbc3d7b9ba", size = 32829770, upload-time = "2026-01-10T21:30:16.359Z" }, + { url = "https://files.pythonhosted.org/packages/48/d1/7b50cedd8c6c9d6f706b4b36fa8544d829c712a75e370f763b318e9638c1/scipy-1.17.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8547e7c57f932e7354a2319fab613981cde910631979f74c9b542bb167a8b9db", size = 35051093, upload-time = "2026-01-10T21:30:22.987Z" }, + { url = "https://files.pythonhosted.org/packages/e2/82/a2d684dfddb87ba1b3ea325df7c3293496ee9accb3a19abe9429bce94755/scipy-1.17.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:33af70d040e8af9d5e7a38b5ed3b772adddd281e3062ff23fec49e49681c38cf", size = 34909905, upload-time = "2026-01-10T21:30:28.704Z" }, + { url = "https://files.pythonhosted.org/packages/ef/5e/e565bd73991d42023eb82bb99e51c5b3d9e2c588ca9d4b3e2cc1d3ca62a6/scipy-1.17.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:f9eb55bb97d00f8b7ab95cb64f873eb0bf54d9446264d9f3609130381233483f", size = 37457743, upload-time = "2026-01-10T21:30:34.819Z" }, + { url = "https://files.pythonhosted.org/packages/58/a8/a66a75c3d8f1fb2b83f66007d6455a06a6f6cf5618c3dc35bc9b69dd096e/scipy-1.17.0-cp314-cp314t-win_amd64.whl", hash = "sha256:1ff269abf702f6c7e67a4b7aad981d42871a11b9dd83c58d2d2ea624efbd1088", size = 37098574, upload-time = "2026-01-10T21:30:40.782Z" }, + { url = "https://files.pythonhosted.org/packages/56/a5/df8f46ef7da168f1bc52cd86e09a9de5c6f19cc1da04454d51b7d4f43408/scipy-1.17.0-cp314-cp314t-win_arm64.whl", hash = "sha256:031121914e295d9791319a1875444d55079885bbae5bdc9c5e0f2ee5f09d34ff", size = 25246266, upload-time = "2026-01-10T21:30:45.923Z" }, +] + +[[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 = "simplejson" +version = "3.20.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/41/f4/a1ac5ed32f7ed9a088d62a59d410d4c204b3b3815722e2ccfb491fa8251b/simplejson-3.20.2.tar.gz", hash = "sha256:5fe7a6ce14d1c300d80d08695b7f7e633de6cd72c80644021874d985b3393649", size = 85784, upload-time = "2025-09-26T16:29:36.64Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/9e/f326d43f6bf47f4e7704a4426c36e044c6bedfd24e072fb8e27589a373a5/simplejson-3.20.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:90d311ba8fcd733a3677e0be21804827226a57144130ba01c3c6a325e887dd86", size = 93530, upload-time = "2025-09-26T16:28:18.07Z" }, + { url = "https://files.pythonhosted.org/packages/35/28/5a4b8f3483fbfb68f3f460bc002cef3a5735ef30950e7c4adce9c8da15c7/simplejson-3.20.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:feed6806f614bdf7f5cb6d0123cb0c1c5f40407ef103aa935cffaa694e2e0c74", size = 75846, upload-time = "2025-09-26T16:28:19.12Z" }, + { url = "https://files.pythonhosted.org/packages/7a/4d/30dfef83b9ac48afae1cf1ab19c2867e27b8d22b5d9f8ca7ce5a0a157d8c/simplejson-3.20.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:6b1d8d7c3e1a205c49e1aee6ba907dcb8ccea83651e6c3e2cb2062f1e52b0726", size = 75661, upload-time = "2025-09-26T16:28:20.219Z" }, + { url = "https://files.pythonhosted.org/packages/09/1d/171009bd35c7099d72ef6afd4bb13527bab469965c968a17d69a203d62a6/simplejson-3.20.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:552f55745044a24c3cb7ec67e54234be56d5d6d0e054f2e4cf4fb3e297429be5", size = 150579, upload-time = "2025-09-26T16:28:21.337Z" }, + { url = "https://files.pythonhosted.org/packages/61/ae/229bbcf90a702adc6bfa476e9f0a37e21d8c58e1059043038797cbe75b8c/simplejson-3.20.2-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c2da97ac65165d66b0570c9e545786f0ac7b5de5854d3711a16cacbcaa8c472d", size = 158797, upload-time = "2025-09-26T16:28:22.53Z" }, + { url = "https://files.pythonhosted.org/packages/90/c5/fefc0ac6b86b9108e302e0af1cf57518f46da0baedd60a12170791d56959/simplejson-3.20.2-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f59a12966daa356bf68927fca5a67bebac0033cd18b96de9c2d426cd11756cd0", size = 148851, upload-time = "2025-09-26T16:28:23.733Z" }, + { url = "https://files.pythonhosted.org/packages/43/f1/b392952200f3393bb06fbc4dd975fc63a6843261705839355560b7264eb2/simplejson-3.20.2-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:133ae2098a8e162c71da97cdab1f383afdd91373b7ff5fe65169b04167da976b", size = 152598, upload-time = "2025-09-26T16:28:24.962Z" }, + { url = "https://files.pythonhosted.org/packages/f4/b4/d6b7279e52a3e9c0fa8c032ce6164e593e8d9cf390698ee981ed0864291b/simplejson-3.20.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7977640af7b7d5e6a852d26622057d428706a550f7f5083e7c4dd010a84d941f", size = 150498, upload-time = "2025-09-26T16:28:26.114Z" }, + { url = "https://files.pythonhosted.org/packages/62/22/ec2490dd859224326d10c2fac1353e8ad5c84121be4837a6dd6638ba4345/simplejson-3.20.2-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:b530ad6d55e71fa9e93e1109cf8182f427a6355848a4ffa09f69cc44e1512522", size = 152129, upload-time = "2025-09-26T16:28:27.552Z" }, + { url = "https://files.pythonhosted.org/packages/33/ce/b60214d013e93dd9e5a705dcb2b88b6c72bada442a97f79828332217f3eb/simplejson-3.20.2-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:bd96a7d981bf64f0e42345584768da4435c05b24fd3c364663f5fbc8fabf82e3", size = 159359, upload-time = "2025-09-26T16:28:28.667Z" }, + { url = "https://files.pythonhosted.org/packages/99/21/603709455827cdf5b9d83abe726343f542491ca8dc6a2528eb08de0cf034/simplejson-3.20.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f28ee755fadb426ba2e464d6fcf25d3f152a05eb6b38e0b4f790352f5540c769", size = 154717, upload-time = "2025-09-26T16:28:30.288Z" }, + { url = "https://files.pythonhosted.org/packages/3c/f9/dc7f7a4bac16cf7eb55a4df03ad93190e11826d2a8950052949d3dfc11e2/simplejson-3.20.2-cp313-cp313-win32.whl", hash = "sha256:472785b52e48e3eed9b78b95e26a256f59bb1ee38339be3075dad799e2e1e661", size = 74289, upload-time = "2025-09-26T16:28:31.809Z" }, + { url = "https://files.pythonhosted.org/packages/87/10/d42ad61230436735c68af1120622b28a782877146a83d714da7b6a2a1c4e/simplejson-3.20.2-cp313-cp313-win_amd64.whl", hash = "sha256:a1a85013eb33e4820286139540accbe2c98d2da894b2dcefd280209db508e608", size = 75972, upload-time = "2025-09-26T16:28:32.883Z" }, + { url = "https://files.pythonhosted.org/packages/05/5b/83e1ff87eb60ca706972f7e02e15c0b33396e7bdbd080069a5d1b53cf0d8/simplejson-3.20.2-py3-none-any.whl", hash = "sha256:3b6bb7fb96efd673eac2e4235200bfffdc2353ad12c54117e1e4e2fc485ac017", size = 57309, upload-time = "2025-09-26T16:29:35.312Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[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 = "tensorboardx" +version = "2.6.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "packaging" }, + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2b/c5/d4cc6e293fb837aaf9f76dd7745476aeba8ef7ef5146c3b3f9ee375fe7a5/tensorboardx-2.6.4.tar.gz", hash = "sha256:b163ccb7798b31100b9f5fa4d6bc22dad362d7065c2f24b51e50731adde86828", size = 4769801, upload-time = "2025-06-10T22:37:07.419Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/1d/b5d63f1a6b824282b57f7b581810d20b7a28ca951f2d5b59f1eb0782c12b/tensorboardx-2.6.4-py3-none-any.whl", hash = "sha256:5970cf3a1f0a6a6e8b180ccf46f3fe832b8a25a70b86e5a237048a7c0beb18e2", size = 87201, upload-time = "2025-06-10T22:37:05.44Z" }, +] + +[[package]] +name = "tensorstore" +version = "0.1.81" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/f6/e2403fc05b97ba74ad408a98a42c288e6e1b8eacc23780c153b0e5166179/tensorstore-0.1.81.tar.gz", hash = "sha256:687546192ea6f6c8ae28d18f13103336f68017d928b9f5a00325e9b0548d9c25", size = 7120819, upload-time = "2026-02-06T18:56:12.535Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/06/c0/32f7d52bfcf1728f557cccb17ac85f57bcc3fa92f4034368d6e7d7d06406/tensorstore-0.1.81-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:7bb563ad4d4d6c4748d9fe4f01f639ddf4ffef83ac180fc3b6d73f46ad854e62", size = 16521316, upload-time = "2026-02-06T18:55:39.557Z" }, + { url = "https://files.pythonhosted.org/packages/38/b9/06ffc44e38ca18aeb3973f6b709d4d2102e17a8d700c7c3e2af3f2830722/tensorstore-0.1.81-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2ff7e6c457596cf21f31c690e451fe634ac804fc98ff8131188e99d5ef7d29bc", size = 14543212, upload-time = "2026-02-06T18:55:42.246Z" }, + { url = "https://files.pythonhosted.org/packages/00/01/3c27962f7258ad0bb552c3cd324fa2e01f746c8b6e81bd25d468f72204e8/tensorstore-0.1.81-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b218a6fe09c72c002f2c6480fc58b78cdbba8bb9c6f3a0d7dd1f70625cb37995", size = 19044489, upload-time = "2026-02-06T18:55:44.957Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ea/fe0f14a1da96d6e0aa6c24d6c31f3ce4b203f8e8a1a2e359489e52b33400/tensorstore-0.1.81-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f33e7c11035c14dad01aeba012051643110cbb95c239e512106fe1be692c98b6", size = 21052658, upload-time = "2026-02-06T18:55:47.138Z" }, + { url = "https://files.pythonhosted.org/packages/e3/e2/cc189d799982f02c200b22405c4d3f28845df6321de2ac3a35ae087758ed/tensorstore-0.1.81-cp313-cp313-win_amd64.whl", hash = "sha256:b55126bcf084cc5fe0151bf465f3a5dedb5b5da0133d01227f75d0e71f9cfae5", size = 13226848, upload-time = "2026-02-06T18:55:49.631Z" }, + { url = "https://files.pythonhosted.org/packages/89/b0/0ca436391f832fad365977623f3c08c4fbbf553fd9a112604aa106646654/tensorstore-0.1.81-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:a48c23e4df50681d8f4f365b08a0beb114ab210accbde9f34d37fd7b45c31005", size = 16525537, upload-time = "2026-02-06T18:55:51.708Z" }, + { url = "https://files.pythonhosted.org/packages/8a/02/c10052b86cf8d47b4cf41e5f139b4003c69bb69e506759b0eb87b873d213/tensorstore-0.1.81-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0be0ce646263820f3d4c9ba738d8e9be7da241cbe093ca2fd02e25023344347c", size = 14547490, upload-time = "2026-02-06T18:55:53.899Z" }, + { url = "https://files.pythonhosted.org/packages/01/d1/bd86c46367624522967e896ca45d77ba9085de3f15081fdad6576ba70aa9/tensorstore-0.1.81-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:93996e756dce82589f5a19e27b4e7c0b5b40221a7e41ddce46dc13d378dbd157", size = 19050938, upload-time = "2026-02-06T18:55:56.123Z" }, + { url = "https://files.pythonhosted.org/packages/11/a2/59a8e9a33cd9e17461f918bda4a20712ed3c51c52e0e42b2f673441bc90d/tensorstore-0.1.81-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:444c088919a739c20ca1f87935d72de4fd87605eb2c0f093b8d49251b7884aef", size = 21055275, upload-time = "2026-02-06T18:55:58.259Z" }, + { url = "https://files.pythonhosted.org/packages/c5/ec/2988f210729b523975b1bee030cabd64b256943c08463331598f1e03bd4f/tensorstore-0.1.81-cp314-cp314-win_amd64.whl", hash = "sha256:f7aa0a3a470c4d832faff7d77dd688b1d352b718d110c95ceba54ec637ca3ffa", size = 13614713, upload-time = "2026-02-06T18:56:00.291Z" }, + { url = "https://files.pythonhosted.org/packages/ae/5d/60e990df3f1dc57c33644375a0eccb906a79fd8a5e2d81238f856c65ad7f/tensorstore-0.1.81-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:6c36d8a827120aa15e50ec5c36dd7e73978d86ba4f46d073fb648d8dda3948e9", size = 16605091, upload-time = "2026-02-06T18:56:02.807Z" }, + { url = "https://files.pythonhosted.org/packages/85/22/f599576815227735d3e34f86f05a8b39d8b15fd979d0029383ebae23978d/tensorstore-0.1.81-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:3c31d831707c4ff3c6ecdcba129f7c39e982572837b2f93e02ccb83fc8581bca", size = 14631573, upload-time = "2026-02-06T18:56:04.892Z" }, + { url = "https://files.pythonhosted.org/packages/cb/76/b5d0b424b7af057a3d4de3f312eba9ddf8a3c750a766b42e0b7f6c2ebef0/tensorstore-0.1.81-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9fba383f108d7450bf9a03487ac7fa3bb2c3080c91cee9d2da3bb217b560846b", size = 19065251, upload-time = "2026-02-06T18:56:06.972Z" }, + { url = "https://files.pythonhosted.org/packages/54/6c/0f113eae73b1e8eb2f712cf5f1efd269452f0f0045158fae43ce7b4701b4/tensorstore-0.1.81-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f88c52f592e2982682045199cabf360462146749d48b7be2969cd640e877c6c3", size = 21066488, upload-time = "2026-02-06T18:56:10.236Z" }, +] + +[[package]] +name = "textual" +version = "7.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py", extra = ["linkify"] }, + { name = "mdit-py-plugins" }, + { name = "platformdirs" }, + { name = "pygments" }, + { name = "rich" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9f/38/7d169a765993efde5095c70a668bf4f5831bb7ac099e932f2783e9b71abf/textual-7.5.0.tar.gz", hash = "sha256:c730cba1e3d704e8f1ca915b6a3af01451e3bca380114baacf6abf87e9dac8b6", size = 1592319, upload-time = "2026-01-30T13:46:39.881Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/78/96ddb99933e11d91bc6e05edae23d2687e44213066bcbaca338898c73c47/textual-7.5.0-py3-none-any.whl", hash = "sha256:849dfee9d705eab3b2d07b33152b7bd74fb1f5056e002873cc448bce500c6374", size = 718164, upload-time = "2026-01-30T13:46:37.635Z" }, +] + +[[package]] +name = "textual-dev" +version = "1.8.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "click" }, + { name = "msgpack" }, + { name = "textual" }, + { name = "textual-serve" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/38/fd/fd5ad9527b536c306a5a860a33a45e13983a59804b48b91ea52b42ba030a/textual_dev-1.8.0.tar.gz", hash = "sha256:7e56867b0341405a95e938cac0647e6d2763d38d0df08710469ad6b6a8db76df", size = 25026, upload-time = "2025-10-11T09:47:01.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/c6/1dc08ceee6d0bdf219c500cb2fbd605b681b63507c4ff27a6477f12f8245/textual_dev-1.8.0-py3-none-any.whl", hash = "sha256:227b6d24a485fbbc77e302aa21f4fdf3083beb57eb45cd95bae082c81cbeddeb", size = 27541, upload-time = "2025-10-11T09:47:00.531Z" }, +] + +[[package]] +name = "textual-serve" +version = "1.1.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "aiohttp-jinja2" }, + { name = "jinja2" }, + { name = "rich" }, + { name = "textual" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c2/7e/62fecc552853ec6a178cb1faa2d6f73b34d5512924770e7b08b58ff14148/textual_serve-1.1.3.tar.gz", hash = "sha256:f8f636ae2f5fd651b79d965473c3e9383d3521cdf896f9bc289709185da3f683", size = 448340, upload-time = "2025-11-01T16:22:36.723Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/fe/108e7773349d500cf363328c3d0b7123e03feda51e310a3a5b136ac8ca71/textual_serve-1.1.3-py3-none-any.whl", hash = "sha256:207a472bc6604e725b1adab4ab8bf12f4c4dc25b04eea31e4d04731d8bf30f18", size = 447339, upload-time = "2025-11-01T16:22:35.209Z" }, +] + +[[package]] +name = "treescope" +version = "0.1.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f0/2a/d13d3c38862632742d2fe2f7ae307c431db06538fd05ca03020d207b5dcc/treescope-0.1.10.tar.gz", hash = "sha256:20f74656f34ab2d8716715013e8163a0da79bdc2554c16d5023172c50d27ea95", size = 138870, upload-time = "2025-08-08T05:43:48.048Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/2b/36e984399089c026a6499ac8f7401d38487cf0183839a4aa78140d373771/treescope-0.1.10-py3-none-any.whl", hash = "sha256:dde52f5314f4c29d22157a6fe4d3bd103f9cae02791c9e672eefa32c9aa1da51", size = 182255, upload-time = "2025-08-08T05:43:46.673Z" }, +] + +[[package]] +name = "types-protobuf" +version = "6.32.1.20251210" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c2/59/c743a842911887cd96d56aa8936522b0cd5f7a7f228c96e81b59fced45be/types_protobuf-6.32.1.20251210.tar.gz", hash = "sha256:c698bb3f020274b1a2798ae09dc773728ce3f75209a35187bd11916ebfde6763", size = 63900, upload-time = "2025-12-10T03:14:25.451Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/aa/43/58e75bac4219cbafee83179505ff44cae3153ec279be0e30583a73b8f108/types_protobuf-6.32.1.20251210-py3-none-any.whl", hash = "sha256:2641f78f3696822a048cfb8d0ff42ccd85c25f12f871fbebe86da63793692140", size = 77921, upload-time = "2025-12-10T03:14:24.477Z" }, +] + +[[package]] +name = "types-requests" +version = "2.32.4.20260107" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0f/f3/a0663907082280664d745929205a89d41dffb29e89a50f753af7d57d0a96/types_requests-2.32.4.20260107.tar.gz", hash = "sha256:018a11ac158f801bfa84857ddec1650750e393df8a004a8a9ae2a9bec6fcb24f", size = 23165, upload-time = "2026-01-07T03:20:54.091Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1c/12/709ea261f2bf91ef0a26a9eed20f2623227a8ed85610c1e54c5805692ecb/types_requests-2.32.4.20260107-py3-none-any.whl", hash = "sha256:b703fe72f8ce5b31ef031264fe9395cac8f46a04661a79f7ed31a80fb308730d", size = 20676, upload-time = "2026-01-07T03:20:52.929Z" }, +] + +[[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-inspect" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mypy-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/74/1789779d91f1961fa9438e9a8710cdae6bd138c80d7303996933d117264a/typing_inspect-0.9.0.tar.gz", hash = "sha256:b23fc42ff6f6ef6954e4852c1fb512cdd18dbea03134f91f856a95ccc9461f78", size = 13825, upload-time = "2023-05-24T20:25:47.612Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/f3/107a22063bf27bdccf2024833d3445f4eea42b2e598abfbd46f6a63b6cb0/typing_inspect-0.9.0-py3-none-any.whl", hash = "sha256:9ee6fc59062311ef8547596ab6b955e1b8aa46242d854bfc78f4f6b0eff35f9f", size = 8827, upload-time = "2023-05-24T20:25:45.287Z" }, +] + +[[package]] +name = "uc-micro-py" +version = "1.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/91/7a/146a99696aee0609e3712f2b44c6274566bc368dfe8375191278045186b8/uc-micro-py-1.0.3.tar.gz", hash = "sha256:d321b92cff673ec58027c04015fcaa8bb1e005478643ff4a500882eaab88c48a", size = 6043, upload-time = "2024-02-09T16:52:01.654Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/87/1f677586e8ac487e29672e4b17455758fce261de06a0d086167bb760361a/uc_micro_py-1.0.3-py3-none-any.whl", hash = "sha256:db1dffff340817673d7b466ec86114a9dc0e9d4d9b5ba229d9d60e5c12600cd5", size = 6229, upload-time = "2024-02-09T16:52:00.371Z" }, +] + +[[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 = "wadler-lindig" +version = "0.1.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1e/67/cbae4bf7683a64755c2c1778c418fea96d00e34395bb91743f08bd951571/wadler_lindig-0.1.7.tar.gz", hash = "sha256:81d14d3fe77d441acf3ebd7f4aefac20c74128bf460e84b512806dccf7b2cd55", size = 15842, upload-time = "2025-06-18T07:00:42.843Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/96/04e7b441807b26b794da5b11e59ed7f83b2cf8af202bd7eba8ad2fa6046e/wadler_lindig-0.1.7-py3-none-any.whl", hash = "sha256:e3ec83835570fd0a9509f969162aeb9c65618f998b1f42918cfc8d45122fe953", size = 20516, upload-time = "2025-06-18T07:00:41.684Z" }, +] + +[[package]] +name = "yarl" +version = "1.22.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "idna" }, + { name = "multidict" }, + { name = "propcache" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/57/63/0c6ebca57330cd313f6102b16dd57ffaf3ec4c83403dcb45dbd15c6f3ea1/yarl-1.22.0.tar.gz", hash = "sha256:bebf8557577d4401ba8bd9ff33906f1376c877aa78d1fe216ad01b4d6745af71", size = 187169, upload-time = "2025-10-06T14:12:55.963Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ea/f3/d67de7260456ee105dc1d162d43a019ecad6b91e2f51809d6cddaa56690e/yarl-1.22.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8dee9c25c74997f6a750cd317b8ca63545169c098faee42c84aa5e506c819b53", size = 139980, upload-time = "2025-10-06T14:10:14.601Z" }, + { url = "https://files.pythonhosted.org/packages/01/88/04d98af0b47e0ef42597b9b28863b9060bb515524da0a65d5f4db160b2d5/yarl-1.22.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:01e73b85a5434f89fc4fe27dcda2aff08ddf35e4d47bbbea3bdcd25321af538a", size = 93424, upload-time = "2025-10-06T14:10:16.115Z" }, + { url = "https://files.pythonhosted.org/packages/18/91/3274b215fd8442a03975ce6bee5fe6aa57a8326b29b9d3d56234a1dca244/yarl-1.22.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:22965c2af250d20c873cdbee8ff958fb809940aeb2e74ba5f20aaf6b7ac8c70c", size = 93821, upload-time = "2025-10-06T14:10:17.993Z" }, + { url = "https://files.pythonhosted.org/packages/61/3a/caf4e25036db0f2da4ca22a353dfeb3c9d3c95d2761ebe9b14df8fc16eb0/yarl-1.22.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b4f15793aa49793ec8d1c708ab7f9eded1aa72edc5174cae703651555ed1b601", size = 373243, upload-time = "2025-10-06T14:10:19.44Z" }, + { url = "https://files.pythonhosted.org/packages/6e/9e/51a77ac7516e8e7803b06e01f74e78649c24ee1021eca3d6a739cb6ea49c/yarl-1.22.0-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5542339dcf2747135c5c85f68680353d5cb9ffd741c0f2e8d832d054d41f35a", size = 342361, upload-time = "2025-10-06T14:10:21.124Z" }, + { url = "https://files.pythonhosted.org/packages/d4/f8/33b92454789dde8407f156c00303e9a891f1f51a0330b0fad7c909f87692/yarl-1.22.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5c401e05ad47a75869c3ab3e35137f8468b846770587e70d71e11de797d113df", size = 387036, upload-time = "2025-10-06T14:10:22.902Z" }, + { url = "https://files.pythonhosted.org/packages/d9/9a/c5db84ea024f76838220280f732970aa4ee154015d7f5c1bfb60a267af6f/yarl-1.22.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:243dda95d901c733f5b59214d28b0120893d91777cb8aa043e6ef059d3cddfe2", size = 397671, upload-time = "2025-10-06T14:10:24.523Z" }, + { url = "https://files.pythonhosted.org/packages/11/c9/cd8538dc2e7727095e0c1d867bad1e40c98f37763e6d995c1939f5fdc7b1/yarl-1.22.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bec03d0d388060058f5d291a813f21c011041938a441c593374da6077fe21b1b", size = 377059, upload-time = "2025-10-06T14:10:26.406Z" }, + { url = "https://files.pythonhosted.org/packages/a1/b9/ab437b261702ced75122ed78a876a6dec0a1b0f5e17a4ac7a9a2482d8abe/yarl-1.22.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:b0748275abb8c1e1e09301ee3cf90c8a99678a4e92e4373705f2a2570d581273", size = 365356, upload-time = "2025-10-06T14:10:28.461Z" }, + { url = "https://files.pythonhosted.org/packages/b2/9d/8e1ae6d1d008a9567877b08f0ce4077a29974c04c062dabdb923ed98e6fe/yarl-1.22.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:47fdb18187e2a4e18fda2c25c05d8251a9e4a521edaed757fef033e7d8498d9a", size = 361331, upload-time = "2025-10-06T14:10:30.541Z" }, + { url = "https://files.pythonhosted.org/packages/ca/5a/09b7be3905962f145b73beb468cdd53db8aa171cf18c80400a54c5b82846/yarl-1.22.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:c7044802eec4524fde550afc28edda0dd5784c4c45f0be151a2d3ba017daca7d", size = 382590, upload-time = "2025-10-06T14:10:33.352Z" }, + { url = "https://files.pythonhosted.org/packages/aa/7f/59ec509abf90eda5048b0bc3e2d7b5099dffdb3e6b127019895ab9d5ef44/yarl-1.22.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:139718f35149ff544caba20fce6e8a2f71f1e39b92c700d8438a0b1d2a631a02", size = 385316, upload-time = "2025-10-06T14:10:35.034Z" }, + { url = "https://files.pythonhosted.org/packages/e5/84/891158426bc8036bfdfd862fabd0e0fa25df4176ec793e447f4b85cf1be4/yarl-1.22.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e1b51bebd221006d3d2f95fbe124b22b247136647ae5dcc8c7acafba66e5ee67", size = 374431, upload-time = "2025-10-06T14:10:37.76Z" }, + { url = "https://files.pythonhosted.org/packages/bb/49/03da1580665baa8bef5e8ed34c6df2c2aca0a2f28bf397ed238cc1bbc6f2/yarl-1.22.0-cp313-cp313-win32.whl", hash = "sha256:d3e32536234a95f513bd374e93d717cf6b2231a791758de6c509e3653f234c95", size = 81555, upload-time = "2025-10-06T14:10:39.649Z" }, + { url = "https://files.pythonhosted.org/packages/9a/ee/450914ae11b419eadd067c6183ae08381cfdfcb9798b90b2b713bbebddda/yarl-1.22.0-cp313-cp313-win_amd64.whl", hash = "sha256:47743b82b76d89a1d20b83e60d5c20314cbd5ba2befc9cda8f28300c4a08ed4d", size = 86965, upload-time = "2025-10-06T14:10:41.313Z" }, + { url = "https://files.pythonhosted.org/packages/98/4d/264a01eae03b6cf629ad69bae94e3b0e5344741e929073678e84bf7a3e3b/yarl-1.22.0-cp313-cp313-win_arm64.whl", hash = "sha256:5d0fcda9608875f7d052eff120c7a5da474a6796fe4d83e152e0e4d42f6d1a9b", size = 81205, upload-time = "2025-10-06T14:10:43.167Z" }, + { url = "https://files.pythonhosted.org/packages/88/fc/6908f062a2f77b5f9f6d69cecb1747260831ff206adcbc5b510aff88df91/yarl-1.22.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:719ae08b6972befcba4310e49edb1161a88cdd331e3a694b84466bd938a6ab10", size = 146209, upload-time = "2025-10-06T14:10:44.643Z" }, + { url = "https://files.pythonhosted.org/packages/65/47/76594ae8eab26210b4867be6f49129861ad33da1f1ebdf7051e98492bf62/yarl-1.22.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:47d8a5c446df1c4db9d21b49619ffdba90e77c89ec6e283f453856c74b50b9e3", size = 95966, upload-time = "2025-10-06T14:10:46.554Z" }, + { url = "https://files.pythonhosted.org/packages/ab/ce/05e9828a49271ba6b5b038b15b3934e996980dd78abdfeb52a04cfb9467e/yarl-1.22.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:cfebc0ac8333520d2d0423cbbe43ae43c8838862ddb898f5ca68565e395516e9", size = 97312, upload-time = "2025-10-06T14:10:48.007Z" }, + { url = "https://files.pythonhosted.org/packages/d1/c5/7dffad5e4f2265b29c9d7ec869c369e4223166e4f9206fc2243ee9eea727/yarl-1.22.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4398557cbf484207df000309235979c79c4356518fd5c99158c7d38203c4da4f", size = 361967, upload-time = "2025-10-06T14:10:49.997Z" }, + { url = "https://files.pythonhosted.org/packages/50/b2/375b933c93a54bff7fc041e1a6ad2c0f6f733ffb0c6e642ce56ee3b39970/yarl-1.22.0-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2ca6fd72a8cd803be290d42f2dec5cdcd5299eeb93c2d929bf060ad9efaf5de0", size = 323949, upload-time = "2025-10-06T14:10:52.004Z" }, + { url = "https://files.pythonhosted.org/packages/66/50/bfc2a29a1d78644c5a7220ce2f304f38248dc94124a326794e677634b6cf/yarl-1.22.0-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ca1f59c4e1ab6e72f0a23c13fca5430f889634166be85dbf1013683e49e3278e", size = 361818, upload-time = "2025-10-06T14:10:54.078Z" }, + { url = "https://files.pythonhosted.org/packages/46/96/f3941a46af7d5d0f0498f86d71275696800ddcdd20426298e572b19b91ff/yarl-1.22.0-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6c5010a52015e7c70f86eb967db0f37f3c8bd503a695a49f8d45700144667708", size = 372626, upload-time = "2025-10-06T14:10:55.767Z" }, + { url = "https://files.pythonhosted.org/packages/c1/42/8b27c83bb875cd89448e42cd627e0fb971fa1675c9ec546393d18826cb50/yarl-1.22.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d7672ecf7557476642c88497c2f8d8542f8e36596e928e9bcba0e42e1e7d71f", size = 341129, upload-time = "2025-10-06T14:10:57.985Z" }, + { url = "https://files.pythonhosted.org/packages/49/36/99ca3122201b382a3cf7cc937b95235b0ac944f7e9f2d5331d50821ed352/yarl-1.22.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:3b7c88eeef021579d600e50363e0b6ee4f7f6f728cd3486b9d0f3ee7b946398d", size = 346776, upload-time = "2025-10-06T14:10:59.633Z" }, + { url = "https://files.pythonhosted.org/packages/85/b4/47328bf996acd01a4c16ef9dcd2f59c969f495073616586f78cd5f2efb99/yarl-1.22.0-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:f4afb5c34f2c6fecdcc182dfcfc6af6cccf1aa923eed4d6a12e9d96904e1a0d8", size = 334879, upload-time = "2025-10-06T14:11:01.454Z" }, + { url = "https://files.pythonhosted.org/packages/c2/ad/b77d7b3f14a4283bffb8e92c6026496f6de49751c2f97d4352242bba3990/yarl-1.22.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:59c189e3e99a59cf8d83cbb31d4db02d66cda5a1a4374e8a012b51255341abf5", size = 350996, upload-time = "2025-10-06T14:11:03.452Z" }, + { url = "https://files.pythonhosted.org/packages/81/c8/06e1d69295792ba54d556f06686cbd6a7ce39c22307100e3fb4a2c0b0a1d/yarl-1.22.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:5a3bf7f62a289fa90f1990422dc8dff5a458469ea71d1624585ec3a4c8d6960f", size = 356047, upload-time = "2025-10-06T14:11:05.115Z" }, + { url = "https://files.pythonhosted.org/packages/4b/b8/4c0e9e9f597074b208d18cef227d83aac36184bfbc6eab204ea55783dbc5/yarl-1.22.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:de6b9a04c606978fdfe72666fa216ffcf2d1a9f6a381058d4378f8d7b1e5de62", size = 342947, upload-time = "2025-10-06T14:11:08.137Z" }, + { url = "https://files.pythonhosted.org/packages/e0/e5/11f140a58bf4c6ad7aca69a892bff0ee638c31bea4206748fc0df4ebcb3a/yarl-1.22.0-cp313-cp313t-win32.whl", hash = "sha256:1834bb90991cc2999f10f97f5f01317f99b143284766d197e43cd5b45eb18d03", size = 86943, upload-time = "2025-10-06T14:11:10.284Z" }, + { url = "https://files.pythonhosted.org/packages/31/74/8b74bae38ed7fe6793d0c15a0c8207bbb819cf287788459e5ed230996cdd/yarl-1.22.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff86011bd159a9d2dfc89c34cfd8aff12875980e3bd6a39ff097887520e60249", size = 93715, upload-time = "2025-10-06T14:11:11.739Z" }, + { url = "https://files.pythonhosted.org/packages/69/66/991858aa4b5892d57aef7ee1ba6b4d01ec3b7eb3060795d34090a3ca3278/yarl-1.22.0-cp313-cp313t-win_arm64.whl", hash = "sha256:7861058d0582b847bc4e3a4a4c46828a410bca738673f35a29ba3ca5db0b473b", size = 83857, upload-time = "2025-10-06T14:11:13.586Z" }, + { url = "https://files.pythonhosted.org/packages/46/b3/e20ef504049f1a1c54a814b4b9bed96d1ac0e0610c3b4da178f87209db05/yarl-1.22.0-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:34b36c2c57124530884d89d50ed2c1478697ad7473efd59cfd479945c95650e4", size = 140520, upload-time = "2025-10-06T14:11:15.465Z" }, + { url = "https://files.pythonhosted.org/packages/e4/04/3532d990fdbab02e5ede063676b5c4260e7f3abea2151099c2aa745acc4c/yarl-1.22.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:0dd9a702591ca2e543631c2a017e4a547e38a5c0f29eece37d9097e04a7ac683", size = 93504, upload-time = "2025-10-06T14:11:17.106Z" }, + { url = "https://files.pythonhosted.org/packages/11/63/ff458113c5c2dac9a9719ac68ee7c947cb621432bcf28c9972b1c0e83938/yarl-1.22.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:594fcab1032e2d2cc3321bb2e51271e7cd2b516c7d9aee780ece81b07ff8244b", size = 94282, upload-time = "2025-10-06T14:11:19.064Z" }, + { url = "https://files.pythonhosted.org/packages/a7/bc/315a56aca762d44a6aaaf7ad253f04d996cb6b27bad34410f82d76ea8038/yarl-1.22.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f3d7a87a78d46a2e3d5b72587ac14b4c16952dd0887dbb051451eceac774411e", size = 372080, upload-time = "2025-10-06T14:11:20.996Z" }, + { url = "https://files.pythonhosted.org/packages/3f/3f/08e9b826ec2e099ea6e7c69a61272f4f6da62cb5b1b63590bb80ca2e4a40/yarl-1.22.0-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:852863707010316c973162e703bddabec35e8757e67fcb8ad58829de1ebc8590", size = 338696, upload-time = "2025-10-06T14:11:22.847Z" }, + { url = "https://files.pythonhosted.org/packages/e3/9f/90360108e3b32bd76789088e99538febfea24a102380ae73827f62073543/yarl-1.22.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:131a085a53bfe839a477c0845acf21efc77457ba2bcf5899618136d64f3303a2", size = 387121, upload-time = "2025-10-06T14:11:24.889Z" }, + { url = "https://files.pythonhosted.org/packages/98/92/ab8d4657bd5b46a38094cfaea498f18bb70ce6b63508fd7e909bd1f93066/yarl-1.22.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:078a8aefd263f4d4f923a9677b942b445a2be970ca24548a8102689a3a8ab8da", size = 394080, upload-time = "2025-10-06T14:11:27.307Z" }, + { url = "https://files.pythonhosted.org/packages/f5/e7/d8c5a7752fef68205296201f8ec2bf718f5c805a7a7e9880576c67600658/yarl-1.22.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bca03b91c323036913993ff5c738d0842fc9c60c4648e5c8d98331526df89784", size = 372661, upload-time = "2025-10-06T14:11:29.387Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2e/f4d26183c8db0bb82d491b072f3127fb8c381a6206a3a56332714b79b751/yarl-1.22.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:68986a61557d37bb90d3051a45b91fa3d5c516d177dfc6dd6f2f436a07ff2b6b", size = 364645, upload-time = "2025-10-06T14:11:31.423Z" }, + { url = "https://files.pythonhosted.org/packages/80/7c/428e5812e6b87cd00ee8e898328a62c95825bf37c7fa87f0b6bb2ad31304/yarl-1.22.0-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:4792b262d585ff0dff6bcb787f8492e40698443ec982a3568c2096433660c694", size = 355361, upload-time = "2025-10-06T14:11:33.055Z" }, + { url = "https://files.pythonhosted.org/packages/ec/2a/249405fd26776f8b13c067378ef4d7dd49c9098d1b6457cdd152a99e96a9/yarl-1.22.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:ebd4549b108d732dba1d4ace67614b9545b21ece30937a63a65dd34efa19732d", size = 381451, upload-time = "2025-10-06T14:11:35.136Z" }, + { url = "https://files.pythonhosted.org/packages/67/a8/fb6b1adbe98cf1e2dd9fad71003d3a63a1bc22459c6e15f5714eb9323b93/yarl-1.22.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:f87ac53513d22240c7d59203f25cc3beac1e574c6cd681bbfd321987b69f95fd", size = 383814, upload-time = "2025-10-06T14:11:37.094Z" }, + { url = "https://files.pythonhosted.org/packages/d9/f9/3aa2c0e480fb73e872ae2814c43bc1e734740bb0d54e8cb2a95925f98131/yarl-1.22.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:22b029f2881599e2f1b06f8f1db2ee63bd309e2293ba2d566e008ba12778b8da", size = 370799, upload-time = "2025-10-06T14:11:38.83Z" }, + { url = "https://files.pythonhosted.org/packages/50/3c/af9dba3b8b5eeb302f36f16f92791f3ea62e3f47763406abf6d5a4a3333b/yarl-1.22.0-cp314-cp314-win32.whl", hash = "sha256:6a635ea45ba4ea8238463b4f7d0e721bad669f80878b7bfd1f89266e2ae63da2", size = 82990, upload-time = "2025-10-06T14:11:40.624Z" }, + { url = "https://files.pythonhosted.org/packages/ac/30/ac3a0c5bdc1d6efd1b41fa24d4897a4329b3b1e98de9449679dd327af4f0/yarl-1.22.0-cp314-cp314-win_amd64.whl", hash = "sha256:0d6e6885777af0f110b0e5d7e5dda8b704efed3894da26220b7f3d887b839a79", size = 88292, upload-time = "2025-10-06T14:11:42.578Z" }, + { url = "https://files.pythonhosted.org/packages/df/0a/227ab4ff5b998a1b7410abc7b46c9b7a26b0ca9e86c34ba4b8d8bc7c63d5/yarl-1.22.0-cp314-cp314-win_arm64.whl", hash = "sha256:8218f4e98d3c10d683584cb40f0424f4b9fd6e95610232dd75e13743b070ee33", size = 82888, upload-time = "2025-10-06T14:11:44.863Z" }, + { url = "https://files.pythonhosted.org/packages/06/5e/a15eb13db90abd87dfbefb9760c0f3f257ac42a5cac7e75dbc23bed97a9f/yarl-1.22.0-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:45c2842ff0e0d1b35a6bf1cd6c690939dacb617a70827f715232b2e0494d55d1", size = 146223, upload-time = "2025-10-06T14:11:46.796Z" }, + { url = "https://files.pythonhosted.org/packages/18/82/9665c61910d4d84f41a5bf6837597c89e665fa88aa4941080704645932a9/yarl-1.22.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:d947071e6ebcf2e2bee8fce76e10faca8f7a14808ca36a910263acaacef08eca", size = 95981, upload-time = "2025-10-06T14:11:48.845Z" }, + { url = "https://files.pythonhosted.org/packages/5d/9a/2f65743589809af4d0a6d3aa749343c4b5f4c380cc24a8e94a3c6625a808/yarl-1.22.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:334b8721303e61b00019474cc103bdac3d7b1f65e91f0bfedeec2d56dfe74b53", size = 97303, upload-time = "2025-10-06T14:11:50.897Z" }, + { url = "https://files.pythonhosted.org/packages/b0/ab/5b13d3e157505c43c3b43b5a776cbf7b24a02bc4cccc40314771197e3508/yarl-1.22.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1e7ce67c34138a058fd092f67d07a72b8e31ff0c9236e751957465a24b28910c", size = 361820, upload-time = "2025-10-06T14:11:52.549Z" }, + { url = "https://files.pythonhosted.org/packages/fb/76/242a5ef4677615cf95330cfc1b4610e78184400699bdda0acb897ef5e49a/yarl-1.22.0-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:d77e1b2c6d04711478cb1c4ab90db07f1609ccf06a287d5607fcd90dc9863acf", size = 323203, upload-time = "2025-10-06T14:11:54.225Z" }, + { url = "https://files.pythonhosted.org/packages/8c/96/475509110d3f0153b43d06164cf4195c64d16999e0c7e2d8a099adcd6907/yarl-1.22.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c4647674b6150d2cae088fc07de2738a84b8bcedebef29802cf0b0a82ab6face", size = 363173, upload-time = "2025-10-06T14:11:56.069Z" }, + { url = "https://files.pythonhosted.org/packages/c9/66/59db471aecfbd559a1fd48aedd954435558cd98c7d0da8b03cc6c140a32c/yarl-1.22.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:efb07073be061c8f79d03d04139a80ba33cbd390ca8f0297aae9cce6411e4c6b", size = 373562, upload-time = "2025-10-06T14:11:58.783Z" }, + { url = "https://files.pythonhosted.org/packages/03/1f/c5d94abc91557384719da10ff166b916107c1b45e4d0423a88457071dd88/yarl-1.22.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e51ac5435758ba97ad69617e13233da53908beccc6cfcd6c34bbed8dcbede486", size = 339828, upload-time = "2025-10-06T14:12:00.686Z" }, + { url = "https://files.pythonhosted.org/packages/5f/97/aa6a143d3afba17b6465733681c70cf175af89f76ec8d9286e08437a7454/yarl-1.22.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:33e32a0dd0c8205efa8e83d04fc9f19313772b78522d1bdc7d9aed706bfd6138", size = 347551, upload-time = "2025-10-06T14:12:02.628Z" }, + { url = "https://files.pythonhosted.org/packages/43/3c/45a2b6d80195959239a7b2a8810506d4eea5487dce61c2a3393e7fc3c52e/yarl-1.22.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:bf4a21e58b9cde0e401e683ebd00f6ed30a06d14e93f7c8fd059f8b6e8f87b6a", size = 334512, upload-time = "2025-10-06T14:12:04.871Z" }, + { url = "https://files.pythonhosted.org/packages/86/a0/c2ab48d74599c7c84cb104ebd799c5813de252bea0f360ffc29d270c2caa/yarl-1.22.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:e4b582bab49ac33c8deb97e058cd67c2c50dac0dd134874106d9c774fd272529", size = 352400, upload-time = "2025-10-06T14:12:06.624Z" }, + { url = "https://files.pythonhosted.org/packages/32/75/f8919b2eafc929567d3d8411f72bdb1a2109c01caaab4ebfa5f8ffadc15b/yarl-1.22.0-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:0b5bcc1a9c4839e7e30b7b30dd47fe5e7e44fb7054ec29b5bb8d526aa1041093", size = 357140, upload-time = "2025-10-06T14:12:08.362Z" }, + { url = "https://files.pythonhosted.org/packages/cf/72/6a85bba382f22cf78add705d8c3731748397d986e197e53ecc7835e76de7/yarl-1.22.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:c0232bce2170103ec23c454e54a57008a9a72b5d1c3105dc2496750da8cfa47c", size = 341473, upload-time = "2025-10-06T14:12:10.994Z" }, + { url = "https://files.pythonhosted.org/packages/35/18/55e6011f7c044dc80b98893060773cefcfdbf60dfefb8cb2f58b9bacbd83/yarl-1.22.0-cp314-cp314t-win32.whl", hash = "sha256:8009b3173bcd637be650922ac455946197d858b3630b6d8787aa9e5c4564533e", size = 89056, upload-time = "2025-10-06T14:12:13.317Z" }, + { url = "https://files.pythonhosted.org/packages/f9/86/0f0dccb6e59a9e7f122c5afd43568b1d31b8ab7dda5f1b01fb5c7025c9a9/yarl-1.22.0-cp314-cp314t-win_amd64.whl", hash = "sha256:9fb17ea16e972c63d25d4a97f016d235c78dd2344820eb35bc034bc32012ee27", size = 96292, upload-time = "2025-10-06T14:12:15.398Z" }, + { url = "https://files.pythonhosted.org/packages/48/b7/503c98092fb3b344a179579f55814b613c1fbb1c23b3ec14a7b008a66a6e/yarl-1.22.0-cp314-cp314t-win_arm64.whl", hash = "sha256:9f6d73c1436b934e3f01df1e1b21ff765cd1d28c77dfb9ace207f746d4610ee1", size = 85171, upload-time = "2025-10-06T14:12:16.935Z" }, + { url = "https://files.pythonhosted.org/packages/73/ae/b48f95715333080afb75a4504487cbe142cae1268afc482d06692d605ae6/yarl-1.22.0-py3-none-any.whl", hash = "sha256:1380560bdba02b6b6c90de54133c81c9f2a453dee9912fe58c1dcced1edb7cff", size = 46814, upload-time = "2025-10-06T14:12:53.872Z" }, +] + +[[package]] +name = "zipp" +version = "3.23.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/02/0f2892c661036d50ede074e376733dca2ae7c6eb617489437771209d4180/zipp-3.23.0.tar.gz", hash = "sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166", size = 25547, upload-time = "2025-06-08T17:06:39.4Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/54/647ade08bf0db230bfea292f893923872fd20be6ac6f53b2b936ba839d75/zipp-3.23.0-py3-none-any.whl", hash = "sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e", size = 10276, upload-time = "2025-06-08T17:06:38.034Z" }, +]