// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "tenzr/modules/checkpoint/checkpoint_manager.h"
#include <algorithm>
#include <chrono>
#include <filesystem>
#include <fstream>
#include <sstream>
#include "tenzr/common/error.h"
#include "tenzr/io/safetensors.h"
namespace tenzr::checkpoint {
CheckpointManager::CheckpointManager(CheckpointOptions options) : options_(std::move(options)) {
if (!std::filesystem::exists(options_.directory)) {
std::filesystem::create_directories(options_.directory);
}
ScanExistingCheckpoints();
}
CheckpointManager::~CheckpointManager() { WaitUntilIdle(); }
CheckpointManager::CheckpointManager(CheckpointManager&& other) noexcept
: options_(std::move(other.options_)),
checkpoints_(std::move(other.checkpoints_)),
best_step_(other.best_step_),
best_metric_(other.best_metric_),
pending_tasks_(std::move(other.pending_tasks_)) {
other.best_step_ = std::nullopt;
other.best_metric_ = 0.0;
}
CheckpointManager& CheckpointManager::operator=(CheckpointManager&& other) noexcept {
if (this != &other) {
WaitUntilIdle();
options_ = std::move(other.options_);
checkpoints_ = std::move(other.checkpoints_);
best_step_ = other.best_step_;
best_metric_ = other.best_metric_;
pending_tasks_ = std::move(other.pending_tasks_);
other.best_step_ = std::nullopt;
other.best_metric_ = 0.0;
}
return *this;
}
void CheckpointManager::ScanExistingCheckpoints() {
checkpoints_.clear();
if (!std::filesystem::exists(options_.directory)) {
return;
}
for (const auto& entry : std::filesystem::directory_iterator(options_.directory)) {
if (entry.is_directory()) {
std::string folder_name = entry.path().filename().string();
if (folder_name.rfind(options_.prefix, 0) == 0) {
std::string step_str = folder_name.substr(options_.prefix.size());
if (step_str.empty() || !std::all_of(step_str.begin(), step_str.end(), [](unsigned char c) { return c >= '0' && c <= '9'; })) {
// Non-numeric folder suffixes are not checkpoints; skip.
continue;
}
std::string model_file = (entry.path() / "model.safetensors").string();
if (std::filesystem::exists(model_file)) {
CheckpointMetadata meta;
meta.step = std::stoull(step_str);
meta.path = entry.path().string();
std::string meta_json = (entry.path() / "meta.json").string();
if (std::filesystem::exists(meta_json)) {
std::ifstream ifs(meta_json);
std::string line;
while (std::getline(ifs, line)) {
if (line.find("\"metric\":") != std::string::npos) {
size_t idx = line.find(':');
#if TENZR_EXCEPTIONS_ENABLED
try {
meta.metric = std::stod(line.substr(idx + 1));
} catch (const std::invalid_argument&) {
meta.metric = 0.0;
} catch (const std::out_of_range&) {
meta.metric = 0.0;
}
#else
(void)idx;
meta.metric = 0.0;
#endif
meta.has_metric = true;
}
}
}
checkpoints_.push_back(meta);
}
}
}
}
std::sort(checkpoints_.begin(), checkpoints_.end(), [](const auto& a, const auto& b) { return a.step < b.step; });
// Determine best step if any
best_step_ = std::nullopt;
for (const auto& cp : checkpoints_) {
if (cp.has_metric) {
if (!best_step_.has_value()) {
best_step_ = cp.step;
best_metric_ = cp.metric;
} else {
bool is_better = options_.higher_is_better ? (cp.metric > best_metric_) : (cp.metric < best_metric_);
if (is_better) {
best_step_ = cp.step;
best_metric_ = cp.metric;
}
}
}
}
}
void CheckpointManager::Save(uint64_t step, const std::unordered_map<std::string, Tensor>& state, std::optional<double> metric) {
if (options_.async) {
// Deep clone state for safe concurrent background serialization
std::unordered_map<std::string, Tensor> cloned_state;
for (const auto& [name, t] : state) {
cloned_state[name] = t.Clone();
}
auto fut = std::async(std::launch::async, [this, step, s = std::move(cloned_state), metric]() mutable { SaveInternal(step, std::move(s), metric); });
pending_tasks_.push_back(std::move(fut));
} else {
SaveInternal(step, state, metric);
}
}
void CheckpointManager::Save(uint64_t step, nn::Module& module, std::optional<double> metric) { Save(step, module.StateDict(), metric); }
void CheckpointManager::SaveInternal(uint64_t step, std::unordered_map<std::string, Tensor> state, std::optional<double> metric) {
auto now = std::chrono::system_clock::now();
uint64_t timestamp = static_cast<uint64_t>(std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count());
std::string tmp_dir =
(std::filesystem::path(options_.directory) / ("tmp_" + options_.prefix + std::to_string(step) + "_" + std::to_string(timestamp))).string();
std::filesystem::create_directories(tmp_dir);
std::string model_path = (std::filesystem::path(tmp_dir) / "model.safetensors").string();
std::unordered_map<std::string, std::string> file_meta;
file_meta["step"] = std::to_string(step);
file_meta["timestamp"] = std::to_string(timestamp);
if (metric.has_value()) {
file_meta["metric"] = std::to_string(*metric);
}
io::SaveSafetensors(model_path, state, file_meta);
// Write meta.json
std::string meta_json = (std::filesystem::path(tmp_dir) / "meta.json").string();
std::ofstream meta_out(meta_json);
meta_out << "{\n";
meta_out << " \"step\": " << step << ",\n";
meta_out << " \"timestamp\": " << timestamp;
if (metric.has_value()) {
meta_out << ",\n \"metric\": " << *metric;
}
meta_out << "\n}\n";
meta_out.close();
// Atomic swap
std::string target_dir = (std::filesystem::path(options_.directory) / (options_.prefix + std::to_string(step))).string();
if (std::filesystem::exists(target_dir)) {
std::filesystem::remove_all(target_dir);
}
std::filesystem::rename(tmp_dir, target_dir);
// Record metadata
CheckpointMetadata cp_meta;
cp_meta.step = step;
cp_meta.path = target_dir;
cp_meta.timestamp = timestamp;
if (metric.has_value()) {
cp_meta.has_metric = true;
cp_meta.metric = *metric;
if (!best_step_.has_value()) {
best_step_ = step;
best_metric_ = *metric;
} else {
bool is_better = options_.higher_is_better ? (*metric > best_metric_) : (*metric < best_metric_);
if (is_better) {
best_step_ = step;
best_metric_ = *metric;
}
}
}
// Add or update checkpoints_
auto it = std::find_if(checkpoints_.begin(), checkpoints_.end(), [step](const auto& cp) { return cp.step == step; });
if (it != checkpoints_.end()) {
*it = cp_meta;
} else {
checkpoints_.push_back(cp_meta);
std::sort(checkpoints_.begin(), checkpoints_.end(), [](const auto& a, const auto& b) { return a.step < b.step; });
}
ApplyRetentionPolicy();
}
void CheckpointManager::ApplyRetentionPolicy() {
if (checkpoints_.size() <= options_.max_to_keep) {
return;
}
// Identify checkpoints to prune from oldest
for (auto it = checkpoints_.begin(); it != checkpoints_.end();) {
if (checkpoints_.size() <= options_.max_to_keep) {
break;
}
if (options_.keep_best && best_step_.has_value() && it->step == *best_step_) {
++it;
continue;
}
if (std::filesystem::exists(it->path)) {
std::filesystem::remove_all(it->path);
}
it = checkpoints_.erase(it);
}
}
void CheckpointManager::WaitUntilIdle() {
for (auto& fut : pending_tasks_) {
if (fut.valid()) {
fut.get();
}
}
pending_tasks_.clear();
}
bool CheckpointManager::HasCheckpoints() const { return !checkpoints_.empty(); }
uint64_t CheckpointManager::LatestStep() const {
TENZR_CHECK(!checkpoints_.empty(), "No checkpoints available in manager");
return checkpoints_.back().step;
}
uint64_t CheckpointManager::BestStep() const {
TENZR_CHECK(best_step_.has_value(), "No best checkpoint recorded");
return *best_step_;
}
std::vector<uint64_t> CheckpointManager::AllSteps() const {
std::vector<uint64_t> steps;
for (const auto& cp : checkpoints_) {
steps.push_back(cp.step);
}
return steps;
}
std::unordered_map<std::string, Tensor> CheckpointManager::Restore(uint64_t step) {
WaitUntilIdle();
std::string target_dir = (std::filesystem::path(options_.directory) / (options_.prefix + std::to_string(step))).string();
std::string model_file = (std::filesystem::path(target_dir) / "model.safetensors").string();
TENZR_CHECK(std::filesystem::exists(model_file), "Checkpoint for step " + std::to_string(step) + " not found: " + model_file);
return io::LoadSafetensors(model_file);
}
std::unordered_map<std::string, Tensor> CheckpointManager::RestoreLatest() { return Restore(LatestStep()); }
std::unordered_map<std::string, Tensor> CheckpointManager::RestoreBest() { return Restore(BestStep()); }
void CheckpointManager::Restore(uint64_t step, nn::Module& module) {
auto state = Restore(step);
module.LoadStateDict(state);
}
void CheckpointManager::RestoreLatest(nn::Module& module) { Restore(LatestStep(), module); }
void CheckpointManager::RestoreBest(nn::Module& module) { Restore(BestStep(), module); }
} // namespace tenzr::checkpoint