Efim Beshmenev
684 lines
24 KiB
C++
684 lines
24 KiB
C++
#include "../../src/TrainingArchive/TrainingArchive.h"
|
|
#include "../../src/TransformerRanker/TransformerRanker.h"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cctype>
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <filesystem>
|
|
#include <iostream>
|
|
#include <limits>
|
|
#include <optional>
|
|
#include <set>
|
|
#include <string>
|
|
#include <tuple>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#ifdef _WIN32
|
|
#include <fcntl.h>
|
|
#include <io.h>
|
|
#include <windows.h>
|
|
#endif
|
|
|
|
namespace {
|
|
|
|
namespace training = szilassi::training;
|
|
namespace transformer = szilassi::transformer;
|
|
|
|
constexpr std::array<char, 8> kExportMagic{{'S', 'Z', 'T', 'X', 'P', '0', '0', '1'}};
|
|
constexpr std::uint32_t kExportVersion = 1;
|
|
constexpr std::size_t kFaceCount = 12;
|
|
constexpr std::size_t kFaceFeatureCount = 16;
|
|
constexpr std::size_t kGlobalFeatureCount = 24;
|
|
constexpr std::size_t kFeatureFloatCount =
|
|
kFaceCount * kFaceFeatureCount + kGlobalFeatureCount;
|
|
|
|
// Stable packed little-endian record layout (there is no native-struct padding):
|
|
// 0 u8[16] canonical UUID bytes
|
|
// 16 u64 record_sequence
|
|
// 24 u32 topology
|
|
// 28 f32[216] proposal features (12x16 face, then 24 global)
|
|
// 892 f32[216] anchor features (12x16 face, then 24 global)
|
|
// 1756 i32 crossing_gain = anchor.C - result.C
|
|
// 1760 i32 intersection_gain = anchor.I - result.I
|
|
// 1764 i32 total_defect_gain
|
|
// 1768 u8 defect_improved (total_defect_gain > 0)
|
|
// 1769 u8 objective_better (TrajectoryImproved)
|
|
// 1770 u8 defect_count_tie (total_defect_gain == 0)
|
|
// 1771 u8 entered_archive
|
|
// 1772 i32[3] plane, move, scale actions (-1 means absent)
|
|
// 1784 f32[3] actual plane, move, scale propensities
|
|
// 1796 u32 replica_count / assigned chains
|
|
// 1800 u64 rollout_iterations
|
|
// 1808 u32 InjectedTrajectory flags
|
|
// 1812 u32 anchor Metrics flags
|
|
// 1816 u32 injected/proposal Metrics flags
|
|
// 1820 u32 result Metrics flags (contains MetricDdVerified when applicable)
|
|
constexpr std::uint32_t kExportRecordSize = 1824;
|
|
|
|
using FeatureFaceArray =
|
|
decltype(std::declval<transformer::Features>().face);
|
|
using FeatureFaceRow = typename FeatureFaceArray::value_type;
|
|
using FeatureGlobalArray =
|
|
decltype(std::declval<transformer::Features>().global);
|
|
static_assert(std::tuple_size<FeatureFaceArray>::value == kFaceCount,
|
|
"Transformer face count changed; bump the export format");
|
|
static_assert(std::tuple_size<FeatureFaceRow>::value == kFaceFeatureCount,
|
|
"Transformer face feature count changed; bump the export format");
|
|
static_assert(std::tuple_size<FeatureGlobalArray>::value == kGlobalFeatureCount,
|
|
"Transformer global feature count changed; bump the export format");
|
|
|
|
struct Options {
|
|
std::filesystem::path root;
|
|
std::set<std::string> included_runs;
|
|
bool root_supplied = false;
|
|
};
|
|
|
|
struct ShardSnapshot {
|
|
std::filesystem::path path;
|
|
std::string run_name;
|
|
training::RunId expected_run_id{};
|
|
std::optional<std::uint64_t> filename_index;
|
|
std::uint64_t size = 0;
|
|
};
|
|
|
|
void print_usage() {
|
|
std::cerr
|
|
<< "Usage: TransformerTrainingExport --root <results/search> "
|
|
"[--include-run <UUID>]...\n"
|
|
<< "Writes SZTXP001 version-1 packed little-endian records to stdout.\n";
|
|
}
|
|
|
|
int hexadecimal_digit(char value) {
|
|
if (value >= '0' && value <= '9') return value - '0';
|
|
if (value >= 'a' && value <= 'f') return 10 + value - 'a';
|
|
if (value >= 'A' && value <= 'F') return 10 + value - 'A';
|
|
return -1;
|
|
}
|
|
|
|
bool parse_canonical_uuid(
|
|
const std::string& text,
|
|
std::string& normalized,
|
|
training::RunId& run_id
|
|
) {
|
|
constexpr std::array<std::size_t, 4> kHyphens{{8, 13, 18, 23}};
|
|
if (text.size() != 36) return false;
|
|
for (std::size_t position : kHyphens) {
|
|
if (text[position] != '-') return false;
|
|
}
|
|
normalized.clear();
|
|
normalized.reserve(text.size());
|
|
std::array<char, 32> hexadecimal{};
|
|
std::size_t digit_count = 0;
|
|
for (std::size_t index = 0; index < text.size(); ++index) {
|
|
if (std::find(kHyphens.begin(), kHyphens.end(), index) != kHyphens.end()) {
|
|
normalized.push_back('-');
|
|
continue;
|
|
}
|
|
const int digit = hexadecimal_digit(text[index]);
|
|
if (digit < 0 || digit_count >= hexadecimal.size()) return false;
|
|
const char lower = static_cast<char>(std::tolower(
|
|
static_cast<unsigned char>(text[index])));
|
|
hexadecimal[digit_count++] = lower;
|
|
normalized.push_back(lower);
|
|
}
|
|
if (digit_count != hexadecimal.size()) return false;
|
|
for (std::size_t index = 0; index < run_id.bytes.size(); ++index) {
|
|
const int high = hexadecimal_digit(hexadecimal[index * 2]);
|
|
const int low = hexadecimal_digit(hexadecimal[index * 2 + 1]);
|
|
run_id.bytes[index] = static_cast<std::uint8_t>((high << 4) | low);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool parse_options(int argc, char* argv[], Options& options) {
|
|
for (int index = 1; index < argc; ++index) {
|
|
const std::string argument = argv[index];
|
|
if (argument == "--help" || argument == "-h") {
|
|
print_usage();
|
|
return false;
|
|
}
|
|
if (argument == "--root") {
|
|
if (++index >= argc || options.root_supplied) {
|
|
std::cerr << "--root requires exactly one path\n";
|
|
return false;
|
|
}
|
|
options.root = std::filesystem::path(argv[index]);
|
|
options.root_supplied = true;
|
|
continue;
|
|
}
|
|
if (argument == "--include-run") {
|
|
if (++index >= argc) {
|
|
std::cerr << "--include-run requires a canonical UUID\n";
|
|
return false;
|
|
}
|
|
std::string normalized;
|
|
training::RunId ignored;
|
|
if (!parse_canonical_uuid(argv[index], normalized, ignored)) {
|
|
std::cerr << "Invalid run UUID: " << argv[index] << '\n';
|
|
return false;
|
|
}
|
|
options.included_runs.insert(std::move(normalized));
|
|
continue;
|
|
}
|
|
std::cerr << "Unknown argument: " << argument << '\n';
|
|
return false;
|
|
}
|
|
if (!options.root_supplied || options.root.empty()) {
|
|
std::cerr << "--root is required\n";
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
std::optional<std::uint64_t> parse_shard_filename(
|
|
const std::filesystem::path& path
|
|
) {
|
|
const std::string filename = path.filename().string();
|
|
constexpr const char* kPrefix = "shard_";
|
|
constexpr const char* kSuffix = ".sztd";
|
|
constexpr std::size_t kPrefixSize = 6;
|
|
constexpr std::size_t kSuffixSize = 5;
|
|
if (filename.size() <= kPrefixSize + kSuffixSize ||
|
|
filename.compare(0, kPrefixSize, kPrefix) != 0 ||
|
|
filename.compare(filename.size() - kSuffixSize, kSuffixSize, kSuffix) != 0) {
|
|
return std::nullopt;
|
|
}
|
|
const std::string digits = filename.substr(
|
|
kPrefixSize,
|
|
filename.size() - kPrefixSize - kSuffixSize);
|
|
std::uint64_t value = 0;
|
|
for (char digit : digits) {
|
|
if (digit < '0' || digit > '9' ||
|
|
value > (std::numeric_limits<std::uint64_t>::max() - 9U) / 10U) {
|
|
return std::nullopt;
|
|
}
|
|
value = value * 10U + static_cast<unsigned>(digit - '0');
|
|
}
|
|
return value;
|
|
}
|
|
|
|
bool snapshot_shards(
|
|
const Options& options,
|
|
std::vector<ShardSnapshot>& shards,
|
|
std::string& error
|
|
) {
|
|
std::error_code path_error;
|
|
const std::filesystem::path absolute_root =
|
|
std::filesystem::absolute(options.root, path_error);
|
|
if (path_error) {
|
|
error = "Cannot resolve archive root: " + path_error.message();
|
|
return false;
|
|
}
|
|
const std::filesystem::path runs_root = absolute_root / "runs";
|
|
if (!std::filesystem::is_directory(runs_root, path_error) || path_error) {
|
|
error = "Archive runs directory does not exist: " + runs_root.string();
|
|
return false;
|
|
}
|
|
|
|
struct RunDirectory {
|
|
std::filesystem::path path;
|
|
std::string normalized_name;
|
|
training::RunId run_id{};
|
|
};
|
|
std::vector<RunDirectory> runs;
|
|
std::set<std::string> found_requested_runs;
|
|
std::filesystem::directory_iterator iterator(runs_root, path_error);
|
|
const std::filesystem::directory_iterator end;
|
|
if (path_error) {
|
|
error = "Cannot enumerate archive runs: " + path_error.message();
|
|
return false;
|
|
}
|
|
while (iterator != end) {
|
|
std::error_code type_error;
|
|
const bool is_directory = iterator->is_directory(type_error);
|
|
if (type_error) {
|
|
error = "Cannot inspect run directory entry: " + type_error.message();
|
|
return false;
|
|
}
|
|
if (is_directory) {
|
|
std::string normalized;
|
|
training::RunId run_id;
|
|
const std::string name = iterator->path().filename().string();
|
|
if (parse_canonical_uuid(name, normalized, run_id) &&
|
|
(options.included_runs.empty() ||
|
|
options.included_runs.count(normalized) != 0)) {
|
|
runs.push_back({iterator->path(), normalized, run_id});
|
|
found_requested_runs.insert(normalized);
|
|
}
|
|
}
|
|
iterator.increment(path_error);
|
|
if (path_error) {
|
|
error = "Cannot continue enumerating archive runs: " + path_error.message();
|
|
return false;
|
|
}
|
|
}
|
|
for (const std::string& requested : options.included_runs) {
|
|
if (found_requested_runs.count(requested) == 0) {
|
|
error = "Requested run does not exist below the archive root: " + requested;
|
|
return false;
|
|
}
|
|
}
|
|
std::sort(runs.begin(), runs.end(), [](const RunDirectory& left, const RunDirectory& right) {
|
|
return left.normalized_name < right.normalized_name;
|
|
});
|
|
|
|
for (const RunDirectory& run : runs) {
|
|
const std::filesystem::path training_directory = run.path / "training";
|
|
std::error_code exists_error;
|
|
const bool has_training =
|
|
std::filesystem::is_directory(training_directory, exists_error);
|
|
if (exists_error) {
|
|
error = "Cannot inspect training directory for run " +
|
|
run.normalized_name + ": " + exists_error.message();
|
|
return false;
|
|
}
|
|
if (!has_training) {
|
|
if (!options.included_runs.empty()) {
|
|
error = "Requested run has no training directory: " + run.normalized_name;
|
|
return false;
|
|
}
|
|
continue;
|
|
}
|
|
std::filesystem::directory_iterator shard_iterator(
|
|
training_directory,
|
|
exists_error);
|
|
if (exists_error) {
|
|
error = "Cannot enumerate training directory for run " +
|
|
run.normalized_name + ": " + exists_error.message();
|
|
return false;
|
|
}
|
|
while (shard_iterator != end) {
|
|
std::error_code file_error;
|
|
const bool regular = shard_iterator->is_regular_file(file_error);
|
|
if (file_error) {
|
|
error = "Cannot inspect training file: " + file_error.message();
|
|
return false;
|
|
}
|
|
const std::filesystem::path path = shard_iterator->path();
|
|
if (regular && path.extension() == ".sztd") {
|
|
const std::uint64_t size = shard_iterator->file_size(file_error);
|
|
if (file_error) {
|
|
error = "Cannot size training shard: " + file_error.message();
|
|
return false;
|
|
}
|
|
shards.push_back({
|
|
path,
|
|
run.normalized_name,
|
|
run.run_id,
|
|
parse_shard_filename(path),
|
|
size});
|
|
}
|
|
shard_iterator.increment(exists_error);
|
|
if (exists_error) {
|
|
error = "Cannot continue enumerating training shards: " +
|
|
exists_error.message();
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
|
|
std::sort(shards.begin(), shards.end(), [](const ShardSnapshot& left,
|
|
const ShardSnapshot& right) {
|
|
if (left.run_name != right.run_name) return left.run_name < right.run_name;
|
|
if (left.filename_index.has_value() != right.filename_index.has_value()) {
|
|
return left.filename_index.has_value();
|
|
}
|
|
if (left.filename_index && right.filename_index &&
|
|
*left.filename_index != *right.filename_index) {
|
|
return *left.filename_index < *right.filename_index;
|
|
}
|
|
return left.path.filename().string() < right.path.filename().string();
|
|
});
|
|
if (shards.empty()) {
|
|
error = "No sealed .sztd shards matched the requested runs";
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
class PackedRecord {
|
|
public:
|
|
void bytes(const std::uint8_t* values, std::size_t count) {
|
|
if (!room(count)) return;
|
|
std::memcpy(data_.data() + size_, values, count);
|
|
size_ += count;
|
|
}
|
|
|
|
void u8(std::uint8_t value) {
|
|
if (!room(1)) return;
|
|
data_[size_++] = value;
|
|
}
|
|
|
|
void u32(std::uint32_t value) {
|
|
if (!room(4)) return;
|
|
for (std::size_t byte = 0; byte < 4; ++byte) {
|
|
data_[size_++] = static_cast<std::uint8_t>(value >> (byte * 8));
|
|
}
|
|
}
|
|
|
|
void i32(std::int32_t value) {
|
|
std::uint32_t bits = 0;
|
|
std::memcpy(&bits, &value, sizeof(bits));
|
|
u32(bits);
|
|
}
|
|
|
|
void u64(std::uint64_t value) {
|
|
if (!room(8)) return;
|
|
for (std::size_t byte = 0; byte < 8; ++byte) {
|
|
data_[size_++] = static_cast<std::uint8_t>(value >> (byte * 8));
|
|
}
|
|
}
|
|
|
|
void f32(float value) {
|
|
std::uint32_t bits = 0;
|
|
static_assert(sizeof(bits) == sizeof(value), "Unexpected float width");
|
|
std::memcpy(&bits, &value, sizeof(bits));
|
|
u32(bits);
|
|
}
|
|
|
|
const std::uint8_t* data() const { return data_.data(); }
|
|
std::size_t size() const { return size_; }
|
|
bool overflowed() const { return overflowed_; }
|
|
|
|
private:
|
|
bool room(std::size_t count) {
|
|
if (count > data_.size() - size_) {
|
|
overflowed_ = true;
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
std::array<std::uint8_t, kExportRecordSize> data_{};
|
|
std::size_t size_ = 0;
|
|
bool overflowed_ = false;
|
|
};
|
|
|
|
bool append_features(
|
|
PackedRecord& packed,
|
|
const transformer::Features& features,
|
|
std::uint32_t topology,
|
|
std::string& error
|
|
) {
|
|
if (features.topology != topology) {
|
|
error = "Transformer feature topology does not match the archive record";
|
|
return false;
|
|
}
|
|
for (const auto& face : features.face) {
|
|
for (float value : face) {
|
|
if (!std::isfinite(value)) {
|
|
error = "Transformer face feature is not finite";
|
|
return false;
|
|
}
|
|
packed.f32(value);
|
|
}
|
|
}
|
|
for (float value : features.global) {
|
|
if (!std::isfinite(value)) {
|
|
error = "Transformer global feature is not finite";
|
|
return false;
|
|
}
|
|
packed.f32(value);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool checked_i32(std::int64_t value, std::int32_t& result) {
|
|
if (value < std::numeric_limits<std::int32_t>::min() ||
|
|
value > std::numeric_limits<std::int32_t>::max()) {
|
|
return false;
|
|
}
|
|
result = static_cast<std::int32_t>(value);
|
|
return true;
|
|
}
|
|
|
|
bool encode_trajectory(
|
|
const training::InjectedTrajectory& trajectory,
|
|
const training::RunId& expected_run_id,
|
|
PackedRecord& packed,
|
|
std::string& error
|
|
) {
|
|
if (trajectory.context.run_id.bytes != expected_run_id.bytes) {
|
|
error = "Training record run UUID does not match its parent run directory";
|
|
return false;
|
|
}
|
|
|
|
const transformer::SearchBudget budget{
|
|
trajectory.rollout_iterations,
|
|
trajectory.replica_count,
|
|
trajectory.replica_count,
|
|
trajectory.pool_size};
|
|
const bool guided =
|
|
(trajectory.flags & training::TrajectoryNeuralGuided) != 0;
|
|
const transformer::FeatureInput proposal_input{
|
|
trajectory.injected_state,
|
|
trajectory.injected_metrics,
|
|
budget,
|
|
trajectory.context.topology,
|
|
guided};
|
|
const transformer::FeatureInput anchor_input{
|
|
trajectory.anchor_state,
|
|
trajectory.anchor_metrics,
|
|
budget,
|
|
trajectory.context.topology,
|
|
guided};
|
|
transformer::Features proposal_features;
|
|
transformer::Features anchor_features;
|
|
std::string feature_error;
|
|
if (!transformer::make_features(
|
|
proposal_input,
|
|
proposal_features,
|
|
&feature_error)) {
|
|
error = "Cannot build proposal features: " + feature_error;
|
|
return false;
|
|
}
|
|
if (!transformer::make_features(
|
|
anchor_input,
|
|
anchor_features,
|
|
&feature_error)) {
|
|
error = "Cannot build anchor features: " + feature_error;
|
|
return false;
|
|
}
|
|
|
|
const std::int64_t crossing_gain_wide =
|
|
static_cast<std::int64_t>(trajectory.anchor_metrics.crossings) -
|
|
static_cast<std::int64_t>(trajectory.result_metrics.crossings);
|
|
const std::int64_t intersection_gain_wide =
|
|
static_cast<std::int64_t>(trajectory.anchor_metrics.intersections) -
|
|
static_cast<std::int64_t>(trajectory.result_metrics.intersections);
|
|
const std::int64_t defect_gain_wide =
|
|
crossing_gain_wide + intersection_gain_wide;
|
|
std::int32_t crossing_gain = 0;
|
|
std::int32_t intersection_gain = 0;
|
|
std::int32_t defect_gain = 0;
|
|
if (!checked_i32(crossing_gain_wide, crossing_gain) ||
|
|
!checked_i32(intersection_gain_wide, intersection_gain) ||
|
|
!checked_i32(defect_gain_wide, defect_gain)) {
|
|
error = "Trajectory label does not fit the export format";
|
|
return false;
|
|
}
|
|
|
|
packed.bytes(trajectory.context.run_id.bytes.data(),
|
|
trajectory.context.run_id.bytes.size());
|
|
packed.u64(trajectory.context.record_sequence);
|
|
packed.u32(trajectory.context.topology);
|
|
if (!append_features(
|
|
packed,
|
|
proposal_features,
|
|
trajectory.context.topology,
|
|
error) ||
|
|
!append_features(
|
|
packed,
|
|
anchor_features,
|
|
trajectory.context.topology,
|
|
error)) {
|
|
return false;
|
|
}
|
|
packed.i32(crossing_gain);
|
|
packed.i32(intersection_gain);
|
|
packed.i32(defect_gain);
|
|
packed.u8(defect_gain > 0 ? 1U : 0U);
|
|
packed.u8((trajectory.flags & training::TrajectoryImproved) != 0 ? 1U : 0U);
|
|
packed.u8(defect_gain == 0 ? 1U : 0U);
|
|
packed.u8((trajectory.flags & training::TrajectoryEnteredArchive) != 0 ? 1U : 0U);
|
|
packed.i32(trajectory.plane_action);
|
|
packed.i32(trajectory.move_action);
|
|
packed.i32(trajectory.scale_action);
|
|
packed.f32(trajectory.plane_propensity);
|
|
packed.f32(trajectory.move_propensity);
|
|
packed.f32(trajectory.scale_propensity);
|
|
packed.u32(trajectory.replica_count);
|
|
packed.u64(trajectory.rollout_iterations);
|
|
packed.u32(trajectory.flags);
|
|
packed.u32(trajectory.anchor_metrics.flags);
|
|
packed.u32(trajectory.injected_metrics.flags);
|
|
packed.u32(trajectory.result_metrics.flags);
|
|
if (packed.overflowed() || packed.size() != kExportRecordSize) {
|
|
error = "Internal export record-size mismatch";
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool write_u32(std::ostream& output, std::uint32_t value) {
|
|
std::array<char, 4> bytes{};
|
|
for (std::size_t index = 0; index < bytes.size(); ++index) {
|
|
bytes[index] = static_cast<char>(
|
|
static_cast<unsigned char>(value >> (index * 8)));
|
|
}
|
|
output.write(bytes.data(), static_cast<std::streamsize>(bytes.size()));
|
|
return static_cast<bool>(output);
|
|
}
|
|
|
|
bool write_header(std::ostream& output) {
|
|
output.write(kExportMagic.data(), static_cast<std::streamsize>(kExportMagic.size()));
|
|
return static_cast<bool>(output) &&
|
|
write_u32(output, kExportVersion) &&
|
|
write_u32(output, kExportRecordSize);
|
|
}
|
|
|
|
void lower_windows_priority() {
|
|
#ifdef _WIN32
|
|
if (SetPriorityClass(GetCurrentProcess(), BELOW_NORMAL_PRIORITY_CLASS) == 0) {
|
|
std::cerr << "Warning: cannot set below-normal process priority (Windows error "
|
|
<< GetLastError() << ")\n";
|
|
}
|
|
#endif
|
|
}
|
|
|
|
bool enable_binary_stdout() {
|
|
#ifdef _WIN32
|
|
return _setmode(_fileno(stdout), _O_BINARY) != -1;
|
|
#else
|
|
return true;
|
|
#endif
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main(int argc, char* argv[]) {
|
|
Options options;
|
|
if (!parse_options(argc, argv, options)) {
|
|
return 2;
|
|
}
|
|
lower_windows_priority();
|
|
if (!enable_binary_stdout()) {
|
|
std::cerr << "Cannot switch stdout to binary mode\n";
|
|
return 3;
|
|
}
|
|
|
|
std::vector<ShardSnapshot> shards;
|
|
std::string error;
|
|
if (!snapshot_shards(options, shards, error)) {
|
|
std::cerr << "Cannot snapshot training shards: " << error << '\n';
|
|
return 4;
|
|
}
|
|
std::uint64_t snapshot_bytes = 0;
|
|
for (const ShardSnapshot& shard : shards) {
|
|
if (shard.size > std::numeric_limits<std::uint64_t>::max() - snapshot_bytes) {
|
|
std::cerr << "Training shard snapshot byte count overflow\n";
|
|
return 4;
|
|
}
|
|
snapshot_bytes += shard.size;
|
|
}
|
|
std::cerr << "Training export snapshot: " << shards.size()
|
|
<< " sealed shards, " << snapshot_bytes << " bytes";
|
|
if (!options.included_runs.empty()) {
|
|
std::cerr << ", " << options.included_runs.size() << " selected run(s)";
|
|
}
|
|
std::cerr << '\n';
|
|
|
|
std::ios::sync_with_stdio(false);
|
|
if (!write_header(std::cout)) {
|
|
std::cerr << "Cannot write training export header to stdout\n";
|
|
return 5;
|
|
}
|
|
|
|
std::uint64_t exported_records = 0;
|
|
for (std::size_t shard_ordinal = 0; shard_ordinal < shards.size(); ++shard_ordinal) {
|
|
const ShardSnapshot& snapshot = shards[shard_ordinal];
|
|
bool output_failed = false;
|
|
std::string record_error;
|
|
training::StreamCallbacks callbacks;
|
|
callbacks.injected_trajectory = [&](const training::InjectedTrajectory& trajectory) {
|
|
PackedRecord packed;
|
|
if (!encode_trajectory(
|
|
trajectory,
|
|
snapshot.expected_run_id,
|
|
packed,
|
|
record_error)) {
|
|
return false;
|
|
}
|
|
std::cout.write(
|
|
reinterpret_cast<const char*>(packed.data()),
|
|
static_cast<std::streamsize>(packed.size()));
|
|
if (!std::cout) {
|
|
output_failed = true;
|
|
return false;
|
|
}
|
|
++exported_records;
|
|
return true;
|
|
};
|
|
training::ShardInfo info;
|
|
std::string shard_error;
|
|
if (!training::stream_read_shard(
|
|
snapshot.path,
|
|
callbacks,
|
|
&info,
|
|
&shard_error)) {
|
|
if (output_failed) {
|
|
std::cerr << "stdout failed while exporting "
|
|
<< snapshot.path.string() << '\n';
|
|
return 5;
|
|
}
|
|
std::cerr << "Cannot export " << snapshot.path.string() << ": "
|
|
<< (record_error.empty() ? shard_error : record_error) << '\n';
|
|
return 6;
|
|
}
|
|
if (snapshot.filename_index && info.shard_index != *snapshot.filename_index) {
|
|
std::cerr << "Shard header index does not match filename: "
|
|
<< snapshot.path.string() << '\n';
|
|
return 6;
|
|
}
|
|
if ((shard_ordinal + 1U) % 16U == 0U ||
|
|
shard_ordinal + 1U == shards.size()) {
|
|
std::cerr << "Validated " << (shard_ordinal + 1U) << '/'
|
|
<< shards.size() << " shards; exported "
|
|
<< exported_records << " trajectories\n";
|
|
}
|
|
}
|
|
std::cout.flush();
|
|
if (!std::cout) {
|
|
std::cerr << "Cannot finish writing the training export stream\n";
|
|
return 5;
|
|
}
|
|
std::cerr << "Training export complete: " << exported_records
|
|
<< " fixed records, record_size=" << kExportRecordSize << '\n';
|
|
return 0;
|
|
}
|