From 09e24221b235dedc40bea72633a9ed525bcfe1c5 Mon Sep 17 00:00:00 2001 From: Giulio Eulisse <10544+ktf@users.noreply.github.com> Date: Mon, 2 Nov 2026 12:05:54 +0100 Subject: [PATCH] DPL Analysis: cap the read buffers of the native AOD reader The native reader fetches the compressed baskets of all the trees a DF is read from with one coalesced read into one buffer, kept until the DF is done, and the buffer pool keeps up to two of them for reuse. On Run 2 AO2Ds (about 62 MB of compressed data per DF) this holds about 75 MB of private memory per reader for the whole job: IgProf shows the largest allocation of the reader, 72 MB, in ReadBufferPool::get, and smaps shows it as anonymous memory, unaffected by fewer threads, no prefetch or MALLOC_ARENA_MAX=1. --aod-reader-buffer-mb N caps that. The trees of a DF, in the order they are read, are packed into read plans of at most N MB of compressed data; a plan is freed when the reads of its trees are done, and the buffer pool keeps at most N MB of free buffers. A DF which fits is still read in one go, so small DFs are not affected, and a single tree larger than N gets a plan of its own. The lookahead for remote files prepares only the trees which fit. Only the trees the reader asks for are read, as before. The default is 0, no limit, which keeps the current behaviour until the default is chosen from measurements. --- .../AODNative/include/AODNative/ByteSource.h | 5 ++ .../AODNative/include/AODNative/ReadPlan.h | 14 +++++- Framework/AODNative/src/ReadPlan.cxx | 42 ++++++++++++++-- Framework/AODNative/test/test_ReadPlan.cxx | 24 +++++++++ .../src/AODJAlienReaderHelpers.cxx | 3 ++ Framework/AnalysisSupport/src/NativeAOD.cxx | 49 +++++++++++++++++-- Framework/AnalysisSupport/src/NativeAOD.h | 24 ++++++--- Framework/Core/src/WorkflowHelpers.cxx | 1 + Framework/Core/src/runDataProcessing.cxx | 1 + 9 files changed, 148 insertions(+), 15 deletions(-) diff --git a/Framework/AODNative/include/AODNative/ByteSource.h b/Framework/AODNative/include/AODNative/ByteSource.h index 113bc3adaef93..26376923328c2 100644 --- a/Framework/AODNative/include/AODNative/ByteSource.h +++ b/Framework/AODNative/include/AODNative/ByteSource.h @@ -36,6 +36,11 @@ struct Options { int threads = 1; int prefetchDepth = 1; size_t coalesceGap = 64 * 1024; + /// Most bytes the read plans of a file hold at once, 0 for no limit. A DF + /// larger than this is read in groups of trees. kAutoBufferBudget: 16 MB for + /// a local file, 64 MB for a remote one, where each read is a round trip. + static constexpr size_t kAutoBufferBudget = SIZE_MAX; + size_t bufferBudget = kAutoBufferBudget; double spanFraction = 0.8; int streamWindows = 4; size_t streamWindowBytes = 2 * 1024 * 1024; diff --git a/Framework/AODNative/include/AODNative/ReadPlan.h b/Framework/AODNative/include/AODNative/ReadPlan.h index b018adbacaaaa..0fbc235244d28 100644 --- a/Framework/AODNative/include/AODNative/ReadPlan.h +++ b/Framework/AODNative/include/AODNative/ReadPlan.h @@ -28,16 +28,26 @@ namespace o2::aod::native class ReadBufferPool { public: + /// Mapped directly rather than taken from malloc, so that a buffer which is + /// dropped always goes back to the system, whatever its size. + struct Unmap { + size_t size = 0; + void operator()(uint8_t* p) const; + }; struct Buffer { - std::unique_ptr data; + std::unique_ptr data; size_t capacity = 0; }; + /// Keeps at most @a maxFreeBytes in free buffers, dropping the largest + /// first; with 0, at most kMaxFree buffers, dropping the smallest first. + explicit ReadBufferPool(size_t maxFreeBytes = 0) : mMaxFreeBytes(maxFreeBytes) {} /// The smallest free buffer of at least @a size bytes, or a new one. Buffer get(size_t size); void put(Buffer buffer); private: static constexpr size_t kMaxFree = 2; + size_t mMaxFreeBytes; std::mutex mMutex; std::vector mFree; }; @@ -58,6 +68,8 @@ class ReadPlan void add(uint64_t offset, size_t length); /// Registers every basket of @a tree which is not embedded in its TBranch. void add(TreeInfo const& tree); + /// The bytes add(@a tree) registers. + static uint64_t bytesOf(TreeInfo const& tree); /// The bytes of [offset, offset + length), if a registered range covers /// them, fetching the pending ranges first if needed. Empty otherwise. The /// span stays valid as long as the plan. diff --git a/Framework/AODNative/src/ReadPlan.cxx b/Framework/AODNative/src/ReadPlan.cxx index 27b8c58695b2b..a093a4504b3bc 100644 --- a/Framework/AODNative/src/ReadPlan.cxx +++ b/Framework/AODNative/src/ReadPlan.cxx @@ -12,6 +12,8 @@ #include "AODNative/ReadPlan.h" #include #include +#include +#include namespace o2::aod::native { @@ -34,17 +36,38 @@ ReadBufferPool::Buffer ReadBufferPool::get(size_t size) // resident next to the new one. Some headroom, so that the next DF, a bit // larger, still fits. mFree.clear(); - auto capacity = size + size / 8; - return {std::make_unique_for_overwrite(capacity), capacity}; + auto capacity = std::max(size + size / 8, 1); + void* p = mmap(nullptr, capacity, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (p == MAP_FAILED) { + throw std::bad_alloc(); + } + return {std::unique_ptr((uint8_t*)p, Unmap{capacity}), capacity}; +} + +void ReadBufferPool::Unmap::operator()(uint8_t* p) const +{ + munmap(p, size); } void ReadBufferPool::put(Buffer buffer) { std::lock_guard lock{mMutex}; mFree.push_back(std::move(buffer)); - if (mFree.size() > kMaxFree) { + if (mMaxFreeBytes == 0 && mFree.size() > kMaxFree) { mFree.erase(std::ranges::min_element(mFree, {}, &Buffer::capacity)); } + auto freeBytes = [this] { + size_t total = 0; + for (auto const& b : mFree) { + total += b.capacity; + } + return total; + }; + // The same headroom as get() allocates, so that a buffer for a group just + // under the budget can still be kept. + while (mMaxFreeBytes > 0 && !mFree.empty() && freeBytes() > mMaxFreeBytes + mMaxFreeBytes / 8) { + mFree.erase(std::ranges::max_element(mFree, {}, &Buffer::capacity)); + } } ReadPlan::ReadPlan(std::shared_ptr source, Options const& options, std::shared_ptr pool) @@ -80,6 +103,19 @@ void ReadPlan::add(TreeInfo const& tree) } } +uint64_t ReadPlan::bytesOf(TreeInfo const& tree) +{ + uint64_t total = 0; + for (auto const& branch : tree.branches) { + for (auto const& basket : branch.baskets) { + if (!basket.embedded) { + total += basket.nbytes; + } + } + } + return total; +} + std::span ReadPlan::lookup(uint64_t offset, size_t length) const { auto it = std::upper_bound(mChunks.begin(), mChunks.end(), offset, [](uint64_t off, Chunk const& c) { return off < c.offset; }); diff --git a/Framework/AODNative/test/test_ReadPlan.cxx b/Framework/AODNative/test/test_ReadPlan.cxx index 2efc1b0d101b8..fc08e28580d8d 100644 --- a/Framework/AODNative/test/test_ReadPlan.cxx +++ b/Framework/AODNative/test/test_ReadPlan.cxx @@ -185,3 +185,27 @@ TEST_CASE("buffers of finished plans are reused by the next one") auto big = pool->get(1 << 16); REQUIRE(big.capacity >= 1 << 16); } + +TEST_CASE("a pool with a budget drops the largest free buffers beyond it") +{ + ReadBufferPool pool{1000}; + auto small = pool.get(400); + auto medium = pool.get(500); + auto large = pool.get(2000); + auto* smallData = small.data.get(); + auto* mediumData = medium.data.get(); + pool.put(std::move(small)); + pool.put(std::move(medium)); + pool.put(std::move(large)); // over the budget: the largest goes + REQUIRE(pool.get(300).data.get() == smallData); + REQUIRE(pool.get(480).data.get() == mediumData); +} + +TEST_CASE("bytesOf counts the baskets a plan reads, not the embedded ones") +{ + TreeInfo tree; + tree.branches.resize(2); + tree.branches[0].baskets = {BasketRef{100, 300, 0, 10, false, {}}, BasketRef{400, 200, 10, 5, false, {}}}; + tree.branches[1].baskets = {BasketRef{600, 50, 0, 10, false, {}}, BasketRef{0, 0, 10, 5, true, {1, 2, 3}}}; + REQUIRE(ReadPlan::bytesOf(tree) == 550); +} diff --git a/Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx b/Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx index 86b7779855bb0..4f0e438651cc2 100644 --- a/Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx +++ b/Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx @@ -206,6 +206,9 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const } context.nativeOptions.threads = std::max(1, options.get("aod-reader-threads")); context.nativeOptions.prefetchDepth = std::max(0, options.get("aod-reader-prefetch")); + if (auto mb = options.get("aod-reader-buffer-mb"); mb >= 0) { + context.nativeOptions.bufferBudget = mb * size_t{1024 * 1024}; + } auto localIO = options.get("aod-reader-local-io"); if (localIO != "mmap" && localIO != "pread") { LOGP(fatal, "Invalid --aod-reader-local-io \"{}\": use mmap or pread.", localIO); diff --git a/Framework/AnalysisSupport/src/NativeAOD.cxx b/Framework/AnalysisSupport/src/NativeAOD.cxx index b5d92015b718f..3d88439c158df 100644 --- a/Framework/AnalysisSupport/src/NativeAOD.cxx +++ b/Framework/AnalysisSupport/src/NativeAOD.cxx @@ -20,8 +20,19 @@ namespace o2::framework using namespace o2::aod::native; +namespace +{ +Options withBudget(Options options, RootFile& file) +{ + if (options.bufferBudget == Options::kAutoBufferBudget) { + options.bufferBudget = (file.source().isRemote() ? 64 : 16) * size_t{1024 * 1024}; + } + return options; +} +} // namespace + NativeAODFileSystem::NativeAODFileSystem(std::shared_ptr file, Options options, RootObjectReadingFactory& factory) - : mFile(std::move(file)), mOptions(options), mObjectFactory(factory) + : mFile(std::move(file)), mOptions(withBudget(options, *mFile)), mObjectFactory(factory), mPool(std::make_shared(mOptions.bufferBudget)) { } @@ -63,12 +74,20 @@ std::shared_ptr NativeAODFileSystem::planFor(std::string const& path) if (df != mPlanDF) { auto previous = std::move(mPlanTrees); mPlanTrees.clear(); + mPlanOf.clear(); + mPlanBytes = 0; mPlanDF = df; // Also when the directory is already known, e.g. the first DF, read at open. adoptLookahead(df); directory(df); if (mAdoptedDF == df) { mPlan = std::move(mAdoptedPlan); + for (auto& name : mAdoptedTrees) { + mPlanOf[name] = mPlan; + if (auto* handle = findTree(df, name)) { + mPlanBytes += ReadPlan::bytesOf(handle->tree); + } + } mPlanTrees = std::move(mAdoptedTrees); } else { mPlan = std::make_shared(mFile->sourcePtr(), mOptions, mPool); @@ -90,7 +109,15 @@ std::shared_ptr NativeAODFileSystem::planFor(std::string const& path) } } addToPlan(treeName); - return mPlan; + // Let go of the plan once its tree is handed out: the plan, and its buffers, + // then go away when the reads of its trees are done. + auto it = mPlanOf.find(treeName); + if (it == mPlanOf.end()) { + return nullptr; + } + auto plan = std::move(it->second); + mPlanOf.erase(it); + return plan; } void NativeAODFileSystem::releaseMappedDF(std::string const& df) @@ -182,10 +209,19 @@ void NativeAODFileSystem::startLookahead(std::string const& afterDF, std::vector Lookahead lookahead{df, {}, {}, std::make_shared(file->sourcePtr(), options, pool)}; file->setTreesToPrefetch(trees); lookahead.keys = file->keysOf(df); + // Only the trees which fit in the budget; the others are read when asked for. + uint64_t bytes = 0; for (auto const& name : trees) { for (auto const& key : lookahead.keys) { if (key.name == name && key.className == "TTree") { - lookahead.trees.emplace_back(name, file->readTree(key)); + auto tree = file->readTree(key); + auto treeBytes = ReadPlan::bytesOf(tree); + if (options.bufferBudget > 0 && bytes > 0 && bytes + treeBytes > options.bufferBudget) { + lookahead.plan->fetch(); + return lookahead; + } + bytes += treeBytes; + lookahead.trees.emplace_back(name, std::move(tree)); lookahead.plan->add(lookahead.trees.back().second); } } @@ -210,7 +246,14 @@ void NativeAODFileSystem::addToPlan(std::string const& treeName) // Reported when the tree itself is asked for. } if (handle) { + auto bytes = ReadPlan::bytesOf(handle->tree); + if (mOptions.bufferBudget > 0 && mPlanBytes > 0 && mPlanBytes + bytes > mOptions.bufferBudget) { + mPlan = std::make_shared(mFile->sourcePtr(), mOptions, mPool); + mPlanBytes = 0; + } mPlan->add(handle->tree); + mPlanBytes += bytes; + mPlanOf[treeName] = mPlan; mPlanTrees.push_back(treeName); } } diff --git a/Framework/AnalysisSupport/src/NativeAOD.h b/Framework/AnalysisSupport/src/NativeAOD.h index a477873875ccb..ec62e6347cff3 100644 --- a/Framework/AnalysisSupport/src/NativeAOD.h +++ b/Framework/AnalysisSupport/src/NativeAOD.h @@ -52,13 +52,17 @@ class NativeAODFileSystem : public VirtualRootFileSystemBase /// The tree at @a path, parsed on first use. nullptr if there is no such /// tree. Throws o2::aod::native::FormatError if it cannot be parsed. NativeTreeHandle* getTreeHandle(std::string const& path); - /// The read plan of the DF of the tree at @a path, with that tree in it; - /// nullptr when the file is mapped and nothing needs reading. A new DF gets a - /// new plan, with the trees read in the previous DF already in it: a reader - /// asks for the same tables in every DF, so the first read of a DF fetches - /// all of them together. For a remote file, the next DF is prepared in the - /// background meanwhile (keys, TTrees and data of the same trees), so that - /// its round trips overlap with the decoding of this one. + /// The read plan with the tree at @a path in it; nullptr when the file is + /// mapped and nothing needs reading. A new DF gets new plans, with the trees + /// read in the previous DF already in them: a reader asks for the same tables + /// in every DF, so the first read of a plan fetches all of its trees + /// together. With Options::bufferBudget, the trees of a DF are split into + /// plans of at most that many bytes, in the order they are read, and a plan + /// is freed once the reads of its trees are done; a tree larger than the + /// budget gets a plan of its own. For a remote file, the next DF is prepared + /// in the background meanwhile (keys, TTrees and data of the same trees, up + /// to the budget), so that its round trips overlap with the decoding of this + /// one. std::shared_ptr planFor(std::string const& path); o2::aod::native::RootFile& file() { return *mFile; } /// The name the file was asked for by (an LFN for a grid file), which is @@ -97,8 +101,12 @@ class NativeAODFileSystem : public VirtualRootFileSystemBase std::mutex mMutex; std::map> mDirectories; std::map> mTrees; - std::shared_ptr mPool = std::make_shared(); + std::shared_ptr mPool; + // The plan trees are being added to, and how many bytes it holds. std::shared_ptr mPlan; + uint64_t mPlanBytes = 0; + // The plan of each tree of the DF, until planFor() hands it out. + std::map> mPlanOf; std::string mPlanDF; // The DF whose pages are mapped, released when the next one starts. std::string mMappedDF; diff --git a/Framework/Core/src/WorkflowHelpers.cxx b/Framework/Core/src/WorkflowHelpers.cxx index c87330f9dbbca..59052efb8bcc1 100644 --- a/Framework/Core/src/WorkflowHelpers.cxx +++ b/Framework/Core/src/WorkflowHelpers.cxx @@ -227,6 +227,7 @@ void WorkflowHelpers::injectServiceDevices(WorkflowSpec& workflow, ConfigContext ConfigParamSpec{"aod-reader-threads", VariantType::Int, 1, {"Decompression threads of the native AOD reader"}}, ConfigParamSpec{"aod-reader-prefetch", VariantType::Int, 1, {"DFs the native AOD reader prepares ahead for remote files: 0 (off) or 1"}}, ConfigParamSpec{"aod-reader-local-io", VariantType::String, "pread", {"How the native AOD reader reads local files: pread or mmap"}}, + ConfigParamSpec{"aod-reader-buffer-mb", VariantType::Int, -1, {"Most MB the native AOD reader holds in read buffers: 0 for no limit, -1 for 16 MB for local files and 64 MB for remote ones"}}, ConfigParamSpec{"aod-reader-json", VariantType::String, {"json configuration file"}}, ConfigParamSpec{"time-limit", VariantType::Int64, 0ll, {"Maximum run time limit in seconds"}}, ConfigParamSpec{"orbit-offset-enumeration", VariantType::Int64, 0ll, {"initial value for the orbit"}}, diff --git a/Framework/Core/src/runDataProcessing.cxx b/Framework/Core/src/runDataProcessing.cxx index c4b2cb892f79d..2877cc58ccf88 100644 --- a/Framework/Core/src/runDataProcessing.cxx +++ b/Framework/Core/src/runDataProcessing.cxx @@ -2089,6 +2089,7 @@ int runStateMachine(DataProcessorSpecs const& workflow, "--aod-reader-threads", "--aod-reader-prefetch", "--aod-reader-local-io", + "--aod-reader-buffer-mb", "--aod-parent-access-level", "--aod-parent-base-path-replacement", "--aod-origin-level-mapping",