PEOutput.h

Go to the documentation of this file.

 
//
// PE node output layout: one file per node (default) or records in one file.
//
// PURPOSE. A DET writes one ZPENode file per branch node, and at long mission
// times those files dominate the run's FILE COUNT (tens of files per
// sequence), which is what a file-count quota limits, not capacity. When such
// a quota is exhausted the run fails with "write failed (disk full or I/O
// error)" from a worker rank even though space is free.
//
// Aggregating the same bytes into one file per rank makes the PE file count
// track the rank count instead of the node count. It is OPT-IN
// (`--aggregate-pe`, default off) because the per-file layout is what the gold
// baselines and the downstream readers (pe-plotter, the clustering tools) read.
//
// RECORD FORMAT. Self-delimiting so that concatenating two aggregate files is
// a valid aggregate file, which lets the MPI merge treat the aggregate as one
// more append file instead of needing rank-unique names:
//
//     #PE <branch_id> <content_bytes>\n
//     <exactly content_bytes of content>\n
//
// The explicit byte count (rather than a sentinel line) keeps parsing
// unambiguous even if node content ever contains a line beginning "#PE ".
// Ids may repeat: MPI clients re-save a whole lineage after every sequence, so
// a later record for an id supersedes an earlier one, exactly as a rewritten
// file supersedes its former contents.
//
// Directory arguments named `pe_dir` are strings that the functions extend by
// plain concatenation, so pass them WITH a trailing '/' (as
// Scheduler::getPEDirectory() returns them). Write failures are not reported
// to the caller, and no function is safe to run concurrently with another one
// on the same directory.
//
#ifndef ADSIDAC_PEOUTPUT_H
#define ADSIDAC_PEOUTPUT_H
 
#include <algorithm>
#include <filesystem>
#include <fstream>
#include <map>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
 
namespace adsidac::pe {
 
namespace fs = std::filesystem;
 
inline const char* aggregate_basename() { return "ZPENodes_agg.txt"; }
 
inline void append_record(const std::string& pe_dir, int branch_id,
                          const std::string& content) {
    std::error_code ec;
    fs::create_directories(pe_dir, ec);
    std::ofstream fout(pe_dir + aggregate_basename(), std::ios::app | std::ios::binary);
    if (!fout) return;
    fout << "#PE " << branch_id << ' ' << content.size() << '\n';
    fout.write(content.data(), static_cast<std::streamsize>(content.size()));
    fout << '\n';
}
 
inline void parse_aggregate(const fs::path& path,
                            std::vector<std::pair<int, std::string>>& out) {
    std::ifstream fin(path, std::ios::binary);
    if (!fin) return;
    std::string line;
    while (std::getline(fin, line)) {
        if (line.rfind("#PE ", 0) != 0) continue;
        std::istringstream hs(line.substr(4));
        int id = 0;
        long long nbytes = -1;
        if (!(hs >> id >> nbytes) || nbytes < 0) break;
        std::string content(static_cast<size_t>(nbytes), '\0');
        fin.read(&content[0], nbytes);
        if (fin.gcount() != nbytes) break;   // truncated final record
        fin.get();                            // trailing newline
        out.emplace_back(id, std::move(content));
    }
}
 
inline std::pair<size_t, size_t> compact_aggregate(const fs::path& agg_path) {
    std::vector<std::pair<int, std::string>> raw;
    parse_aggregate(agg_path, raw);
    if (raw.empty()) return {0, 0};
    std::map<int, std::string> latest;
    for (auto& [id, content] : raw)
        latest[id] = std::move(content);
    if (latest.size() == raw.size()) return {raw.size(), latest.size()};
    const fs::path tmp = agg_path.string() + ".compact";
    {
        std::ofstream fout(tmp, std::ios::binary | std::ios::trunc);
        if (!fout) return {raw.size(), latest.size()};
        for (auto& [id, content] : latest) {
            fout << "#PE " << id << ' ' << content.size() << '\n';
            fout.write(content.data(), static_cast<std::streamsize>(content.size()));
            fout << '\n';
        }
    }
    std::error_code ec;
    fs::rename(tmp, agg_path, ec);
    if (ec) fs::remove(tmp, ec);
    return {raw.size(), latest.size()};
}
 
inline std::vector<std::pair<int, std::string>> collect(const std::string& pe_dir) {
    std::vector<std::pair<int, std::string>> raw;
    std::error_code ec;
    if (!fs::exists(pe_dir, ec)) return {};
    for (auto& entry : fs::recursive_directory_iterator(pe_dir, ec)) {
        if (ec) break;
        if (!entry.is_regular_file()) continue;
        const std::string fname = entry.path().filename().string();
        if (fname == aggregate_basename()) {
            parse_aggregate(entry.path(), raw);
        } else if (fname.size() > 7 && fname.compare(0, 7, "ZPENode") == 0
                   && fname.back() == 't') {
            const size_t us = fname.find('_'), dot = fname.find('.');
            if (us == std::string::npos || dot == std::string::npos || dot <= us) continue;
            int id = 0;
            try { id = std::stoi(fname.substr(us + 1, dot - us - 1)); }
            catch (...) { continue; }
            std::ifstream fin(entry.path());
            std::string content((std::istreambuf_iterator<char>(fin)),
                                std::istreambuf_iterator<char>());
            raw.emplace_back(id, std::move(content));
        }
    }
    std::map<int, std::string> latest;   // last write wins, ordered by id
    for (auto& [id, content] : raw) {
        if (id < 0) continue;
        latest[id] = std::move(content);
    }
    return {latest.begin(), latest.end()};
}
 
inline size_t rewrite_sequential(const std::string& pe_dir,
                                 std::vector<std::pair<int, std::string>> entries,
                                 bool aggregate) {
    std::sort(entries.begin(), entries.end(),
              [](const auto& a, const auto& b) { return a.first < b.first; });
    std::error_code ec;
    fs::remove_all(pe_dir, ec);
    const std::string folder = pe_dir + "PEFolder_0/";
    fs::create_directories(folder, ec);
    if (aggregate) {
        std::ofstream fout(folder + aggregate_basename(), std::ios::binary);
        for (size_t i = 0; i < entries.size(); ++i) {
            fout << "#PE " << i << ' ' << entries[i].second.size() << '\n';
            fout.write(entries[i].second.data(),
                       static_cast<std::streamsize>(entries[i].second.size()));
            fout << '\n';
        }
    } else {
        for (size_t i = 0; i < entries.size(); ++i) {
            std::ofstream fout(folder + "ZPENode_" + std::to_string(i) + ".txt");
            fout << entries[i].second;
        }
    }
    return entries.size();
}
 
inline size_t finalize_sequential(const std::string& pe_dir, bool aggregate) {
    std::error_code ec;
    if (!fs::exists(pe_dir, ec)) return 0;
    if (aggregate)
        return rewrite_sequential(pe_dir, collect(pe_dir), true);
 
    const std::string journal_path = pe_dir + "/renumber.map";
    const std::string folder = (fs::path(pe_dir) / "PEFolder_0").string() + "/";
    std::vector<std::pair<int, int>> mapping;   // old id -> new id
 
    std::ifstream jin(journal_path);
    if (jin) {
        // Crash recovery: resume the journaled mapping, never recompute it.
        std::string header;
        std::getline(jin, header);
        int o = 0, n = 0;
        while (jin >> o >> n) mapping.emplace_back(o, n);
        if (header.rfind("# adsidac pe renumber v1", 0) != 0 || mapping.empty())
            return rewrite_sequential(pe_dir, collect(pe_dir), false);
    } else {
        // Fresh scan: filenames only. Any aggregate file or duplicate id means
        // mixed/unexpected state - take the content path instead.
        std::map<int, fs::path> ids;
        for (auto& entry : fs::recursive_directory_iterator(pe_dir, ec)) {
            if (ec) break;
            if (!entry.is_regular_file()) continue;
            const std::string fname = entry.path().filename().string();
            if (fname == aggregate_basename())
                return rewrite_sequential(pe_dir, collect(pe_dir), false);
            if (fname.size() > 7 && fname.compare(0, 7, "ZPENode") == 0
                && fname.back() == 't') {
                const size_t us = fname.find('_'), dot = fname.find('.');
                if (us == std::string::npos || dot == std::string::npos || dot <= us) continue;
                int id = 0;
                try { id = std::stoi(fname.substr(us + 1, dot - us - 1)); }
                catch (...) { continue; }
                if (id < 0) continue;   // sentinels removed in the cleanup pass
                if (!ids.emplace(id, entry.path()).second)
                    return rewrite_sequential(pe_dir, collect(pe_dir), false);
            }
        }
        if (ids.empty()) return 0;
        // Idempotency: already finalized (dense 0..N-1, everything in PEFolder_0).
        bool done = ids.rbegin()->first == static_cast<int>(ids.size()) - 1;
        if (done)
            for (auto& [id, p] : ids)
                if (p.parent_path().filename() != "PEFolder_0") { done = false; break; }
        if (!done) {
            int next = 0;
            for (auto& [id, p] : ids) mapping.emplace_back(id, next++);
            std::ofstream jout(journal_path, std::ios::trunc);
            jout << "# adsidac pe renumber v1 " << mapping.size() << '\n';
            for (auto& [o, n] : mapping) jout << o << ' ' << n << '\n';
            jout.flush();
            if (!jout) return rewrite_sequential(pe_dir, collect(pe_dir), false);
        }
    }
 
    if (!mapping.empty()) {
        fs::create_directories(folder, ec);
        // Pass 1: old name -> tmp_<new>. Existence checks make recovery re-runs
        // and fresh runs the same loop.
        for (auto& [oldid, newid] : mapping) {
            const std::string final_name = folder + "ZPENode_" + std::to_string(newid) + ".txt";
            const std::string tmp_name = folder + "tmp_" + std::to_string(newid);
            if (fs::exists(final_name, ec) || fs::exists(tmp_name, ec)) continue;
            const std::string old_name = (fs::path(pe_dir) / ("PEFolder_" + std::to_string(oldid / 1000))
                                          / ("ZPENode_" + std::to_string(oldid) + ".txt")).string();
            if (fs::exists(old_name, ec)) {
                fs::rename(old_name, tmp_name, ec);
            } else {
                // Bucket layouts can differ (a serial resume writes epoch ids into
                // their own buckets); find the file wherever it is.
                for (auto& entry : fs::recursive_directory_iterator(pe_dir, ec)) {
                    if (!entry.is_regular_file()) continue;
                    if (entry.path().filename() == "ZPENode_" + std::to_string(oldid) + ".txt") {
                        fs::rename(entry.path(), tmp_name, ec);
                        break;
                    }
                }
            }
        }
        // Pass 2: tmp_<new> -> ZPENode_<new>.txt.
        for (auto& [oldid, newid] : mapping) {
            const std::string tmp_name = folder + "tmp_" + std::to_string(newid);
            if (fs::exists(tmp_name, ec))
                fs::rename(tmp_name, folder + "ZPENode_" + std::to_string(newid) + ".txt", ec);
        }
    }
 
    // Cleanup: negative-id sentinels, emptied buckets, the journal. The
    // content path gets the same result from rewrite_sequential's remove_all;
    // the finalized-run detector relies on the sentinels being gone.
    size_t count = 0;
    std::vector<fs::path> doomed;
    for (auto& entry : fs::recursive_directory_iterator(pe_dir, ec)) {
        if (!entry.is_regular_file()) continue;
        const std::string fname = entry.path().filename().string();
        if (fname.size() > 8 && fname.compare(0, 8, "ZPENode_") == 0 && fname[8] == '-')
            doomed.push_back(entry.path());
        else if (fname.compare(0, 7, "ZPENode") == 0 && fname.back() == 't')
            ++count;
    }
    for (auto& p : doomed) fs::remove(p, ec);
    for (auto& entry : fs::directory_iterator(pe_dir, ec)) {
        if (entry.is_directory() && entry.path().filename() != "PEFolder_0"
            && fs::is_empty(entry.path(), ec))
            fs::remove(entry.path(), ec);
    }
    fs::remove(journal_path, ec);
    return count;
}
 
} // namespace adsidac::pe
 
#endif // ADSIDAC_PEOUTPUT_H