diff --git a/core/imt/CMakeLists.txt b/core/imt/CMakeLists.txt index c99e8cd53ece9..e042fab222151 100644 --- a/core/imt/CMakeLists.txt +++ b/core/imt/CMakeLists.txt @@ -32,6 +32,7 @@ if(imt) ROOT/RSlotStack.hxx ROOT/TExecutor.hxx ROOT/TThreadExecutor.hxx + ROOT/InternalIMTUtils.hxx LINKDEF LinkDef.h MODULE diff --git a/core/imt/inc/ROOT/InternalIMTUtils.hxx b/core/imt/inc/ROOT/InternalIMTUtils.hxx new file mode 100644 index 0000000000000..ca2feabf88862 --- /dev/null +++ b/core/imt/inc/ROOT/InternalIMTUtils.hxx @@ -0,0 +1,42 @@ +#ifndef ROOT_INTERNAL_IMTUTILS +#define ROOT_INTERNAL_IMTUTILS + +#include +#include + +namespace ROOT { +class TThreadExecutor; +} + +namespace ROOT::Internal::IMTUtils { + +/// \brief Abstract implementation of the oneAPI Range requirements +/// +/// For the specification, see +/// https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range. +/// +/// The class is meant to pass implementations from other parts of ROOT to the +/// internal calls to the TBB parallel_for overload taking a Range object. +class RAbstractTBBRangeImpl { +public: + virtual std::unique_ptr SplitWork() = 0; + virtual bool Empty() const = 0; + virtual bool IsDivisible() const = 0; + RAbstractTBBRangeImpl() = default; + + // Destructor is required for dynamic polymorphism, so apply rule of five + virtual ~RAbstractTBBRangeImpl() = default; + RAbstractTBBRangeImpl(const RAbstractTBBRangeImpl &) = default; + RAbstractTBBRangeImpl &operator=(const RAbstractTBBRangeImpl &) = default; + RAbstractTBBRangeImpl(RAbstractTBBRangeImpl &&) = default; + RAbstractTBBRangeImpl &operator=(RAbstractTBBRangeImpl &&) = default; +}; + +// Entry point to a parallel_for that takes a Range and Body +// https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/algorithms/functions/parallel_for_func +void ParallelFor(ROOT::TThreadExecutor &pool, std::unique_ptr rangeImpl, + const std::function &body); + +} // namespace ROOT::Internal::IMTUtils + +#endif diff --git a/core/imt/inc/ROOT/TThreadExecutor.hxx b/core/imt/inc/ROOT/TThreadExecutor.hxx index dd1494f0c1cf1..11ceb0f506c82 100644 --- a/core/imt/inc/ROOT/TThreadExecutor.hxx +++ b/core/imt/inc/ROOT/TThreadExecutor.hxx @@ -27,6 +27,7 @@ #include "ROOT/TypeTraits.hxx" // InvokeResult #include "RTaskArena.hxx" #include "TError.h" +#include #include //std::function #include @@ -40,6 +41,9 @@ namespace ROOT { class TThreadExecutor: public TExecutorCRTP { friend TExecutorCRTP; + friend void ROOT::Internal::IMTUtils::ParallelFor( + TThreadExecutor &pool, std::unique_ptr rangeImpl, + const std::function &body); public: @@ -124,6 +128,8 @@ namespace ROOT { // Functions that interface with the parallel library used as a backend void ParallelFor(unsigned start, unsigned end, unsigned step, const std::function &f); + void ParallelFor(std::unique_ptr rangeImpl, + const std::function &body); double ParallelReduce(const std::vector &objs, const std::function &redfunc); float ParallelReduce(const std::vector &objs, const std::function &redfunc); template diff --git a/core/imt/src/TThreadExecutor.cxx b/core/imt/src/TThreadExecutor.cxx index 6fd443cf6b66d..0becf7c652b45 100644 --- a/core/imt/src/TThreadExecutor.cxx +++ b/core/imt/src/TThreadExecutor.cxx @@ -214,3 +214,53 @@ unsigned TThreadExecutor::GetPoolSize() const } } // namespace ROOT + +namespace { +/// \brief The Range used in the TThreadExecutor::ParallelFor call below +/// +/// A simple struct in an anonymous namespace to be able to use the tbb::split tag in the basic splitting constructor. +/// https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range +/// +/// The struct implements the Range requirements: +/// - Copy constructor and destructor (implicitly done by using rule of zero) +/// - "Basic splitting constructor", a constructor taking tbb::split tag which calls the SplitWork method of the +/// implementation to create a new task with roughly half the work of the original one. +/// - is_divisible to check if the task can be divided in two. +/// - empty to check if the task is empty. +struct TThreadExecutorTBBRange { + + std::shared_ptr fImpl; + + // Constructor to capture implementation from outside + TThreadExecutorTBBRange(std::unique_ptr impl) + : fImpl(std::move(impl)) + { + } + // Basic splitting constructor + TThreadExecutorTBBRange(TThreadExecutorTBBRange &r, tbb::split) : fImpl(r.fImpl->SplitWork()) {} + // True if range can be partitioned into two subranges + bool is_divisible() const { return fImpl->IsDivisible(); } + // True if range is empty + bool empty() const { return fImpl->Empty(); } +}; + +} // namespace + +void ROOT::TThreadExecutor::ParallelFor( + std::unique_ptr rangeImpl, + const std::function &body) +{ + fTaskArenaW->Access().execute([&] { + tbb::this_task_arena::isolate([&] { + tbb::parallel_for(TThreadExecutorTBBRange{std::move(rangeImpl)}, + [body](const TThreadExecutorTBBRange &r) { body(*r.fImpl); }); + }); + }); +} + +void ROOT::Internal::IMTUtils::ParallelFor( + ROOT::TThreadExecutor &pool, std::unique_ptr rangeImpl, + const std::function &body) +{ + pool.ParallelFor(std::move(rangeImpl), body); +} diff --git a/tree/dataframe/CMakeLists.txt b/tree/dataframe/CMakeLists.txt index 506151d24ab81..c172a1324ebb2 100644 --- a/tree/dataframe/CMakeLists.txt +++ b/tree/dataframe/CMakeLists.txt @@ -183,4 +183,8 @@ target_include_directories(ROOTDataFrame PRIVATE ../../builtins/exprtk) # hard-c target_link_libraries(ROOTDataFrame PRIVATE nlohmann_json::nlohmann_json) +if(imt) + target_link_libraries(ROOTDataFrame PRIVATE TBB::tbb) +endif() + ROOT_ADD_TEST_SUBDIRECTORY(test) diff --git a/tree/dataframe/doc/parallel_scheduling.md b/tree/dataframe/doc/parallel_scheduling.md new file mode 100644 index 0000000000000..a73f4d4acb1be --- /dev/null +++ b/tree/dataframe/doc/parallel_scheduling.md @@ -0,0 +1,67 @@ +# Parallel scheduling in RDataFrame + +RDataFrame allows to process input data in parallel in two ways: with implicit multi-threading and with distributed multiprocessing. This document explains the strategies for the implementations of the two modes. + +[//]: # (This should be a good way to write comments in markdown according to https://stackoverflow.com/questions/4823468/comments-in-markdown) +[//]: # (TODO: The document only features a description of the multithreading strategy for RNTuple, add the rest.) + +## Multi-threading + +### RNTuple + +RNTuple data reading is handled in `RNTupleDS`, an `RDataSource`-derived class. The reading is triggered by the +computation graph being run, and nodes requesting data from the data source. The computation graph is started and +managed by an `RLoopManager`, which delegates to the `RNTupleDS` for decisions related to scheduling of different MT +tasks, reading data and giving it back to the `RLoopManager` to run the computation graph until done. This is a +description of the high-level steps involved: + +1. When the user triggers the computation graph, `RLoopManager::RunDataSourceMT` prepares the computation graph nodes + and calls `RDataSource::ProcessMT`. This method is overridden by the RNTuple data source. This method, as well as all + other methods described from now on, are private to the `RNTupleDS` class. +2. In `RNTupleDS::ProcessMT`: + a. A `TThreadExecutor` is initialised with the number of threads given by `ROOT::GetThreadPoolSize()`. + b. An `RSlotStack` is initialised with the same number of slots. A slot is required to apply the computation graph + to a range of entries. + c. Optionally, if the user has specified a global entry range to restrict processing to, a list of global entry + offsets per file is populated by getting the number of entries in each file upfront. + d. One file-level task per input file is scheduled in a `tbb::task_group`. The task group is executed inside the + arena managed by `TThreadExecutor`. +3. The `RNTupleDS` constructor opens the first input file, reads its schema and leaves the corresponding page source in + the staging area. `RNTupleDS::Initialize` starts a background staging thread. The staging thread opens files in + order, calls `RPageSource::LoadStructure` and puts the page sources into the staging area. The sliding window of + files that are staged ahead of the next file to be claimed is bounded by the number of slots. `fNextFileToClaim` is + an atomic counter advanced by the file-level tasks; the staging thread only stages files while + `fNextFileToStage < min(nFiles, fNextFileToClaim + nSlots)`. A condition variable wakes the staging thread when a + file is claimed and wakes file-level tasks when their page source is ready. +4. A file-level task atomically claims the next input file. If it claims the last file, it sets the shared tail-mode + flag. The task then takes the staged page source, calls `RPageSource::Attach` and + `RPageSource::LoadAllPageLists`, and computes the list of clusters to process. With a global entry range, the + cluster boundaries are trimmed to the intersection of the file and the requested range; clusters that do not + contain any requested entry are dropped. +5. The file-level task acquires one processing slot and iterates over its clusters in file order. Before starting, it + restricts the page source entry range to the whole file (or to the trimmed range when a global entry range is + active), not to the current cluster. This is important: it lets the page source prefetch the next clusters while the + computation graph processes the current one. For every cluster, the task inserts the active entry range for its + slot and calls `RLoopManager::RNTupleThreadTask`. The entry range passed to the loop manager is expressed in global + entry coordinates by adding the file offset. +6. Tail scheduling starts once every file-level task has been claimed. This is the point at which no new file can + start and some workers may otherwise be left without work while the remaining clusters of a large file are still + being processed. A file-level task that observes tail mode schedules one task for each of its remaining clusters in + the same `tbb::task_group` and then processes its current cluster. This happens at most once per file-level task. +7. A tail cluster task first acquires a processing slot and only then clones the parent page source. It restricts the + clone to the single cluster it is going to process and calls `RLoopManager::RNTupleThreadTask`, exactly like the + file-level task does for a cluster. Cloning after acquiring the slot is what keeps the number of concurrently open + page sources bounded. + +The staging window, the file-level tasks and the tail cluster tasks together bound the number of open page sources. +At any point there are at most `nSlots` page sources in the staging area and at most `nSlots` page sources owned by +file-level tasks or their tail cluster tasks. Tail cluster tasks are only created by file-level tasks that have +already claimed a file, and there are at most `nSlots` such file-level tasks active at once. Consequently, the data +source keeps at most `2 * nSlots` page sources open. This property holds also during the transition to tail mode and +for `nSlots == 1`. + +This strategy deliberately avoids the eager, recursive splitting of the previous implementation. Splitting a file +into ranges as soon as the computation starts destroys the page source's ability to prefetch across cluster +boundaries and can produce tasks that are too small to amortise their setup cost. With the current strategy, a file +is read sequentially for as long as there is other file-level work to claim, and cluster-level parallelism is only +introduced for the tail of the event loop. diff --git a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx index 6b6af72f84331..47e5e747ed76f 100644 --- a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx @@ -330,12 +330,19 @@ public: return fSuppressErrorsForMissingBranches; } + // TODO(vpadulan): it seems undesirable that we need different "thread tasks" depending on the data source. It would + // be nice to merge everything into one, potentially opening up the possibility to making the MT processing interface + // of the data source public. + /// The task run by every thread on the input entry range, for the generic RDataSource. void DataSourceThreadTask(const std::pair &entryRange, ROOT::Internal::RSlotStack &slotStack, std::atomic &entryCount); /// The task run by every thread on an entry range (known by the input TTreeReader), for the TTree data source. void TTreeThreadTask(TTreeReader &treeReader, ROOT::Internal::RSlotStack &slotStack, std::atomic &entryCount); + /// The task run by every thread on the input entry range, for the RNTuple data source. + void RNTupleThreadTask(const std::pair &entryRange, unsigned int slot, + std::uint64_t columnReaderOffset); }; /// \brief Create an RLoopManager that reads a TChain. diff --git a/tree/dataframe/inc/ROOT/RNTupleDS.hxx b/tree/dataframe/inc/ROOT/RNTupleDS.hxx index 8bddba044e01c..9faf397bf8d5e 100644 --- a/tree/dataframe/inc/ROOT/RNTupleDS.hxx +++ b/tree/dataframe/inc/ROOT/RNTupleDS.hxx @@ -21,8 +21,10 @@ #include #include +#include #include #include +#include #include #include #include @@ -35,6 +37,36 @@ namespace ROOT { class RDataFrame; } + +namespace ROOT::Detail::RDF { +class RLoopManager; +} + +namespace ROOT::Internal { +class RSlotStack; +} + +namespace ROOT { +class RFieldBase; +class RDataFrame; +class RNTuple; +} // namespace ROOT +namespace ROOT::Detail::RDF { +class RNodeBase; +} +namespace ROOT::RDF { +template +class RInterface; +} +namespace ROOT::Internal::RDF { +class RNTupleColumnReader; +std::vector> +GetDatasetGlobalClusterBoundaries(const ROOT::RDF::RInterface &node); +} // namespace ROOT::Internal::RDF +namespace ROOT::Internal { +class RPageSource; +} + namespace ROOT::Internal::RDF { /** * \brief Internal overload of the function that allows passing a range of entries @@ -57,39 +89,40 @@ std::pair, ROOT::NTupleSiz GetClustersAndEntries(std::string_view ntupleName, std::string_view location); } // namespace ROOT::Internal::RDF +#ifdef R__USE_IMT namespace ROOT { -class RFieldBase; -class RDataFrame; -class RNTuple; -} // namespace ROOT -namespace ROOT::Detail::RDF { -class RNodeBase; -} -namespace ROOT::RDF { -template -class RInterface; -} -namespace ROOT::Internal::RDF { -class RNTupleColumnReader; -std::vector> -GetDatasetGlobalClusterBoundaries(const ROOT::RDF::RInterface &node); -} -namespace ROOT::Internal { -class RPageSource; +class TThreadExecutor; } +#endif namespace ROOT::RDF { class RNTupleDS final : public ROOT::RDF::RDataSource { friend class ROOT::Internal::RDF::RNTupleColumnReader; - /// The PrepareNextRanges() method populates the fNextRanges list with REntryRangeDS records. - /// The GetEntryRanges() swaps fNextRanges and fCurrentRanges and uses the list of - /// REntryRangeDS records to return the list of ranges ready to use by the RDF loop manager. + /// \brief Page source information consumed while running the computation graph + /// + /// REntryRangeDS objects are populated differently in single-thread and multi-thread mode. + /// + /// In single-thread mode, the PrepareNextRanges() method populates the fNextRanges list with REntryRangeDS records. + /// The GetEntryRanges() swaps fNextRanges and fCurrentRanges and uses the list of REntryRangeDS records to return + /// the list of ranges ready to use by the RDF loop manager. + /// + /// In multi-thread mode, the task being executed by a thread injects an REntryRangeDS into the fActiveRangesPerSlot + /// member, which is then retrievable for relevant operations e.g. connecting the column readers to the page source + /// and populating the sample information. struct REntryRangeDS { - std::unique_ptr fSource; + /// The page source. It can be shared between this object and other multithreaded tasks. + std::shared_ptr fSource; ULong64_t fFirstEntry = 0; ///< First entry index in fSource /// End entry index in fSource, e.g. the number of entries in the range is fLastEntry - fFirstEntry ULong64_t fLastEntry = 0; + + // TODO(vpadulan): Having to have the entry offset is ugly, currently it's necessary because the implementation of + // RNTupleColumnReader::GetImpl interprets the 'entry' input argument as a "global" entry w.r.t. some offset. For + // lack of a better communication strategy at the moment, we carry this offset over with REntryRangeDS as it is + // also used by the RLoopManager e.g. to decide the values for rdfentry_ for a given slot + + ULong64_t fEntryOffset = 0; /// Offset of both first and last entries w.r.t. the page source std::string_view fFileName; ///< Storage location of the current RNTuple }; @@ -115,7 +148,7 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { /// and /// c) trigger staging of the next batch of files in the I/O background thread. /// 4. On `Finalize()`, the I/O background thread is stopped. - std::vector> fStagingArea; + std::vector> fStagingArea; std::size_t fNextFileIndex = 0; ///< Index into fFileNames to the next file to process /// We prepare a prototype field for every column. If a column reader is actually requested @@ -143,14 +176,15 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { std::vector fCurrentRanges; ///< Basis for the ranges returned by the last GetEntryRanges() call std::vector fNextRanges; ///< Basis for the ranges populated by the PrepareNextRanges() call + + /// The page source and entry window seen by an active thread task in multi-threaded runs + std::vector fActiveRangesPerSlot; /// Maps the first entries from the ranges of the last GetEntryRanges() call to their corresponding index in /// the fCurrentRanges vectors. This is necessary because the returned ranges get distributed arbitrarily /// onto slots. In the InitSlot method, the column readers use this map to find the correct range to connect to. std::unordered_map fFirstEntry2RangeIdx; // Keep track of the scheduled entries - necessary for processing of GlobalEntries std::vector> fOriginalRanges; - /// One element per slot, corresponding to the current range index for that slot, as filled by InitSlot - std::vector fSlotsToRangeIdxs; /// The background thread that runs StageNextSources() std::thread fThreadStaging; @@ -164,6 +198,16 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { bool fHasNextSources = false; /// Is true when the I/O thread should quit bool fStagingThreadShouldTerminate = false; +#ifdef R__USE_IMT + /// True when the current event loop uses the multi-threaded scheduler + bool fIsMT = false; + /// Index of the next file to be claimed by a file-level MT task + std::atomic fNextFileToClaim{0}; + /// Index of the next file to be staged by the background MT staging thread + std::size_t fNextFileToStage = 0; + /// Exception raised by the background MT staging thread, if any + std::exception_ptr fStagingError; +#endif /// \brief Holds useful information about fields added to the RNTupleDS struct RFieldInfo { @@ -201,6 +245,13 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { /// Starting from `fNextFileIndex`, opens the next `fNSlots` files. Calls `LoadStructure()` on the opened files. /// The very first file is already available from the constructor. void StageNextSources(); +#ifdef R__USE_IMT + /// The main function of the background staging thread used by the multi-threaded scheduler. + /// It stages files ahead of the next file to be claimed, keeping at most fNSlots files in the window. + void ExecStagingMT(); + /// Waits until the page source for fileIdx is staged and moves it out of the staging area. + std::shared_ptr TakeStagedPageSource(std::size_t fileIdx); +#endif /// Populates fNextRanges with the next set of entry ranges. Moves files from the staging area as necessary /// and aligns ranges with cluster boundaries for scheduling the tail of files. /// Upon return, the fNextRanges list is ordered. It has usually fNSlots elements; fewer if there @@ -221,7 +272,17 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { explicit RNTupleDS(std::string_view ntupleName, const std::vector &fileNames, const std::pair &range); +#ifdef R__USE_IMT + // Helpers for the MT scheduling + + void InsertActiveEntryRange(unsigned int slot, std::string_view fileName, + std::shared_ptr pageSource, std::uint64_t beginEntry, + std::uint64_t endEntry, std::uint64_t offset); + std::vector ComputeGlobalEntryOffsets(ROOT::TThreadExecutor &pool) const; + void ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) final; + +#endif public: RNTupleDS(std::string_view ntupleName, std::string_view fileName); RNTupleDS(std::string_view ntupleName, const std::vector &fileNames); diff --git a/tree/dataframe/src/RLoopManager.cxx b/tree/dataframe/src/RLoopManager.cxx index 591b3790666e9..d750b4dfa8ddf 100644 --- a/tree/dataframe/src/RLoopManager.cxx +++ b/tree/dataframe/src/RLoopManager.cxx @@ -1343,6 +1343,40 @@ void ROOT::Detail::RDF::RLoopManager::DataSourceThreadTask(const std::pair &entryRange, + unsigned int slot, std::uint64_t columnReaderOffset) +{ +#ifdef R__USE_IMT + // These are begin and end entries of the current cluster in the file currently opened by the thread who owns the + // slot, offset by a global atomic counter with the value passed via columnReaderOffset. The only reason we use it + // for the moment is to provide a seed for a unique rdfentry_ sequence in the current slot task. + const auto [begin, end] = entryRange; + + RDSRangeRAII _{*this, slot, columnReaderOffset}; + RCallCleanUpTask cleanup(*this, slot); + + fSampleInfos[slot] = ROOT::Internal::RDF::CreateSampleInfo(*fDataSource, slot, fSampleMap); + + R__LOG_DEBUG(0, RDFLogChannel()) << LogRangeProcessing( + {fDataSource->GetLabel(), begin - columnReaderOffset, end - columnReaderOffset, slot}); + + try { + for (auto entry = begin; entry < end; ++entry) { + if (fDataSource->SetEntry(slot, entry)) { + RunAndCheckFilters(slot, entry); + } + } + } catch (...) { + std::cerr << "RDataFrame::Run: event loop was interrupted\n"; + throw; + } +#else + (void)entryRange; + (void)slot; + (void)columnReaderOffset; +#endif +} + void ROOT::Detail::RDF::RLoopManager::TTreeThreadTask(TTreeReader &treeReader, ROOT::Internal::RSlotStack &slotStack, std::atomic &entryCount) { diff --git a/tree/dataframe/src/RNTupleDS.cxx b/tree/dataframe/src/RNTupleDS.cxx index 135dd04f18ac8..4c43ae1fed8bb 100644 --- a/tree/dataframe/src/RNTupleDS.cxx +++ b/tree/dataframe/src/RNTupleDS.cxx @@ -2,6 +2,8 @@ /// \author Jakob Blomer /// \author Enrico Guiraud /// \date 2018-10-04 +/// \author Vincenzo Eduardo Padulano CERN +/// \date 2026-10-19 /************************************************************************* * Copyright (C) 1995-2020, Rene Brun and Fons Rademakers. * @@ -27,6 +29,7 @@ #include #include +#include #include #include #include @@ -37,6 +40,11 @@ #include #include #include +#ifdef R__USE_IMT +#include +#include +#include +#endif // clang-format off /** @@ -303,6 +311,8 @@ ROOT::RDF::RNTupleDS::RNTupleDS(std::unique_ptr pag { pageSource->Attach(); fPrincipalDescriptor = pageSource->GetSharedDescriptorGuard()->Clone(); + // The first file is used to read the schema and is also the first entry in the staging area. Keeping it + // avoids reopening the file at the start of every event loop. fStagingArea.emplace_back(std::move(pageSource)); AddField(fPrincipalDescriptor, "", fPrincipalDescriptor.GetFieldZeroId(), @@ -502,6 +512,66 @@ void ROOT::RDF::RNTupleDS::StageNextSources() } } +#ifdef R__USE_IMT +void ROOT::RDF::RNTupleDS::ExecStagingMT() +{ + const auto nFiles = fFileNames.empty() ? std::size_t{1} : fFileNames.size(); + std::unique_lock lock(fMutexStaging); + + while (!fStagingThreadShouldTerminate) { + const auto nextFile = fNextFileToClaim.load(std::memory_order_acquire); + const auto windowEnd = std::min(nFiles, nextFile + fNSlots); + + if (fNextFileToStage >= nFiles) + return; + + if (fNextFileToStage >= windowEnd) { + fCvStaging.wait(lock, [&] { + return fStagingThreadShouldTerminate || + fNextFileToStage < std::min(nFiles, fNextFileToClaim.load(std::memory_order_acquire) + fNSlots); + }); + continue; + } + + const auto fileIdx = fNextFileToStage++; + lock.unlock(); + try { + if (fFileNames.empty()) + throw std::runtime_error("RNTupleDS: cannot stage a file without a file name"); + auto source = CreatePageSource(fNTupleName, fFileNames[fileIdx]); + source->LoadStructure(); + lock.lock(); + fStagingArea[fileIdx] = std::move(source); + } catch (...) { + lock.lock(); + fStagingError = std::current_exception(); + fCvStaging.notify_all(); + return; + } + fCvStaging.notify_all(); + } +} + +std::shared_ptr ROOT::RDF::RNTupleDS::TakeStagedPageSource(std::size_t fileIdx) +{ + std::unique_lock lock(fMutexStaging); + fCvStaging.wait(lock, [&] { + return fStagingArea[fileIdx] || fStagingError || fStagingThreadShouldTerminate; + }); + + if (fStagingError) + std::rethrow_exception(fStagingError); + if (!fStagingArea[fileIdx]) + throw std::runtime_error("RNTupleDS: background staging stopped before the requested file was available"); + + auto source = std::move(fStagingArea[fileIdx]); + lock.unlock(); + // Wake the staging thread so that it can advance its sliding window. + fCvStaging.notify_all(); + return source; +} +#endif + void ROOT::RDF::RNTupleDS::PrepareNextRanges() { assert(fNextRanges.empty()); @@ -537,70 +607,11 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges() } return; } - - // Work scheduling of the tail: multiple slots work on the same file. - // Every slot still has its own page source but these page sources may open the same file. - // Again, we need to skip empty files. - unsigned int nSlotsPerFile = fNSlots / nRemainingFiles; - for (std::size_t i = 0; (fNextRanges.size() < fNSlots) && (fNextFileIndex < nFiles); ++i) { - std::unique_ptr source; - // Need to look for the file name to populate the sample info later - const auto &sourceFileName = fFileNames[fNextFileIndex]; - std::swap(fStagingArea[fNextFileIndex], source); - if (!source) { - // Empty files trigger this condition - source = CreatePageSource(fNTupleName, fFileNames[fNextFileIndex]); - } - source->Attach(); - source->LoadAllPageLists(); - fNextFileIndex++; - - auto nEntries = source->GetNEntries(); - if (nEntries == 0) - continue; - - // If last file: use all remaining slots - if (i == (nRemainingFiles - 1)) - nSlotsPerFile = fNSlots - fNextRanges.size(); - - const auto rangesByCluster = [&source]() { - // Take the shared lock of the descriptor just for the time necessary - auto descGuard = source->GetSharedDescriptorGuard(); - return ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()); - }(); - - const unsigned int nRangesByCluster = rangesByCluster.size(); - - // Distribute slots equidistantly over the entry range, aligned on cluster boundaries - const auto nClustersPerSlot = nRangesByCluster / nSlotsPerFile; - const auto remainder = nRangesByCluster % nSlotsPerFile; - std::size_t iRange = 0; - unsigned int iSlot = 0; - const unsigned int N = std::min(nSlotsPerFile, nRangesByCluster); - for (; iSlot < N; ++iSlot) { - auto start = rangesByCluster[iRange].fFirstEntry; - iRange += nClustersPerSlot + static_cast(iSlot < remainder); - assert(iRange > 0); - auto end = rangesByCluster[iRange - 1].fLastEntryPlusOne; - - REntryRangeDS range; - range.fFileName = sourceFileName; - // The last range for this file just takes the already opened page source. All previous ranges clone. - if (iSlot == N - 1) { - range.fSource = std::move(source); - } else { - range.fSource = source->Clone(); - } - range.fSource->SetEntryRange({start, end - start}); - range.fFirstEntry = start; - range.fLastEntry = end; - fNextRanges.emplace_back(std::move(range)); - } - } // loop over tail of remaining files } std::vector> ROOT::RDF::RNTupleDS::GetEntryRanges() { + assert(!ROOT::IsImplicitMTEnabled()); // MT scheduling doesn't use this function std::vector> ranges; // We need to distinguish between single threaded and multi-threaded runs. @@ -608,11 +619,8 @@ std::vector> ROOT::RDF::RNTupleDS::GetEntryRange // to new page sources of the chain in GetEntryRanges. In multi-threaded mode, on the other hand, // InitSlot is called for every returned range, thus rewiring the column readers takes place in // InitSlot and FinalizeSlot. - - if (fNSlots == 1) { - for (auto r : fActiveColumnReaders[0]) { - r->Disconnect(true /* keepValue */); - } + for (auto r : fActiveColumnReaders[0]) { + r->Disconnect(true /* keepValue */); } // If we have fewer files than slots and we run multiple event loops, we can reuse fCurrentRanges and don't need @@ -724,7 +732,7 @@ std::vector> ROOT::RDF::RNTupleDS::GetEntryRange fSeenEntriesNoGlobalRange += nEntriesPerSource; - if ((fNSlots == 1) && (fCurrentRanges[0].fSource)) { + if (fCurrentRanges[0].fSource) { for (auto r : fActiveColumnReaders[0]) { r->Connect(*fCurrentRanges[0].fSource, fOriginalRanges[0].first); } @@ -733,32 +741,24 @@ std::vector> ROOT::RDF::RNTupleDS::GetEntryRange return ranges; } -void ROOT::RDF::RNTupleDS::InitSlot(unsigned int slot, ULong64_t firstEntry) +void ROOT::RDF::RNTupleDS::InitSlot(unsigned int slot, ULong64_t columnReaderOffset) { - if (fNSlots == 1) { - // Ensure the connection between slot and range is valid also in single-thread mode - fSlotsToRangeIdxs[0] = 0; + // Only MT scheduling uses this function + if (!ROOT::IsImplicitMTEnabled()) return; - } - // The same slot ID could be picked multiple times in the same execution, thus - // ending up processing different page sources. Here we re-establish the - // connection between the slot and the correct page source by finding which - // range index corresponds to the first entry passed. - auto idxRange = fFirstEntry2RangeIdx.at(firstEntry); - - // We also remember this connection so it can later be retrieved in CreateSampleInfo - fSlotsToRangeIdxs[slot * ROOT::Internal::RDF::CacheLineStep()] = idxRange; + const auto &entryRangeDS = fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()); + assert(entryRangeDS.fSource != nullptr); for (auto r : fActiveColumnReaders[slot]) { - r->Connect(*fCurrentRanges[idxRange].fSource, - fOriginalRanges[idxRange].first - fCurrentRanges[idxRange].fFirstEntry); + r->Connect(*entryRangeDS.fSource, columnReaderOffset); } } void ROOT::RDF::RNTupleDS::FinalizeSlot(unsigned int slot) { - if (fNSlots == 1) + // Only MT scheduling uses this function + if (!ROOT::IsImplicitMTEnabled()) return; for (auto r : fActiveColumnReaders[slot]) { @@ -789,9 +789,27 @@ void ROOT::RDF::RNTupleDS::Initialize() fSeenEntriesNoGlobalRange = 0; fNextFileIndex = 0; fIsReadyForStaging = fHasNextSources = fStagingThreadShouldTerminate = false; - fThreadStaging = std::thread(&RNTupleDS::ExecStaging, this); assert(fNextRanges.empty()); +#ifdef R__USE_IMT + fIsMT = ROOT::IsImplicitMTEnabled(); + if (fIsMT) { + const auto nFiles = fFileNames.empty() ? std::size_t{1} : fFileNames.size(); + if (fStagingArea.size() != nFiles) + fStagingArea.resize(nFiles); + { + std::lock_guard lock(fMutexStaging); + fNextFileToClaim.store(0, std::memory_order_relaxed); + fNextFileToStage = fStagingArea[0] ? 1 : 0; + fStagingError = nullptr; + } + fThreadStaging = std::thread(&RNTupleDS::ExecStagingMT, this); + return; + } +#endif + + fThreadStaging = std::thread(&RNTupleDS::ExecStaging, this); + if (fCurrentRanges.empty() || (fFileNames.size() > fNSlots)) { // First event loop or large number of files: start the staging process. { @@ -817,8 +835,31 @@ void ROOT::RDF::RNTupleDS::Finalize() std::lock_guard _(fMutexStaging); fStagingThreadShouldTerminate = true; } - fCvStaging.notify_one(); - fThreadStaging.join(); + fCvStaging.notify_all(); + +#ifdef R__USE_IMT + if (fIsMT) { + if (fThreadStaging.joinable()) + fThreadStaging.join(); + + std::lock_guard lock(fMutexStaging); + // Drop all page sources opened by the staging thread. The next event loop starts with a clean + // staging area and re-opens the first file in the background as needed. + fStagingArea.clear(); + fStagingArea.resize(fFileNames.empty() ? std::size_t{1} : fFileNames.size()); + fNextFileToClaim.store(0, std::memory_order_relaxed); + fNextFileToStage = 0; + fStagingError = nullptr; + for (auto &range : fActiveRangesPerSlot) + range = {}; + fCurrentRanges.clear(); + fNextRanges.clear(); + return; + } +#endif + + if (fThreadStaging.joinable()) + fThreadStaging.join(); // If we have a chain with more files than the number of slots, the files opened at the end of the // event loop won't be reused when the event loop restarts, so we can close them. if (fFileNames.size() > fNSlots) { @@ -835,7 +876,7 @@ void ROOT::RDF::RNTupleDS::SetNSlots(unsigned int nSlots) assert(nSlots > 0); fNSlots = nSlots; fActiveColumnReaders.resize(fNSlots); - fSlotsToRangeIdxs.resize(fNSlots * ROOT::Internal::RDF::CacheLineStep()); + fActiveRangesPerSlot.resize(fNSlots * ROOT::Internal::RDF::CacheLineStep()); } ROOT::RDataFrame ROOT::RDF::FromRNTuple(std::string_view ntupleName, std::string_view fileName) @@ -856,27 +897,31 @@ ROOT::RDF::RSampleInfo ROOT::Internal::RDF::RNTupleDS::CreateSampleInfo( // connection between the slot and the correct page source by retrieving // which range is connected currently to the slot - const auto &rangeIdx = fSlotsToRangeIdxs.at(slot * ROOT::Internal::RDF::CacheLineStep()); + const auto &entryRangeDS = (!ROOT::IsImplicitMTEnabled()) + ? fCurrentRanges[0] + : fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()); // Missing source if a file does not exist - if (!fCurrentRanges[rangeIdx].fSource) + if (!entryRangeDS.fSource) return ROOT::RDF::RSampleInfo{}; - const auto &ntupleName = fCurrentRanges[rangeIdx].fSource->GetNTupleName(); - const auto &ntuplePath = fCurrentRanges[rangeIdx].fFileName; + const auto &ntupleName = entryRangeDS.fSource->GetNTupleName(); + const auto &ntuplePath = entryRangeDS.fFileName; const auto ntupleID = std::string(ntuplePath) + '/' + ntupleName; if (sampleMap.empty()) - return ROOT::RDF::RSampleInfo( - ntupleID, std::make_pair(fCurrentRanges[rangeIdx].fFirstEntry, fCurrentRanges[rangeIdx].fLastEntry), nullptr, - fPrincipalDescriptor.GetNEntries()); + return ROOT::RDF::RSampleInfo(ntupleID, + std::make_pair(entryRangeDS.fFirstEntry - entryRangeDS.fEntryOffset, + entryRangeDS.fLastEntry - entryRangeDS.fEntryOffset), + nullptr, fPrincipalDescriptor.GetNEntries()); if (sampleMap.find(ntupleID) == sampleMap.end()) throw std::runtime_error("Full sample identifier '" + ntupleID + "' cannot be found in the available samples."); - return ROOT::RDF::RSampleInfo( - ntupleID, std::make_pair(fCurrentRanges[rangeIdx].fFirstEntry, fCurrentRanges[rangeIdx].fLastEntry), - sampleMap.at(ntupleID), fPrincipalDescriptor.GetNEntries()); + return ROOT::RDF::RSampleInfo(ntupleID, + std::make_pair(entryRangeDS.fFirstEntry - entryRangeDS.fEntryOffset, + entryRangeDS.fLastEntry - entryRangeDS.fEntryOffset), + sampleMap.at(ntupleID), fPrincipalDescriptor.GetNEntries()); } ROOT::RDataFrame ROOT::Internal::RDF::FromRNTuple(std::string_view ntupleName, @@ -897,3 +942,232 @@ ROOT::Internal::RDF::GetClustersAndEntries(std::string_view ntupleName, std::str auto descGuard = source->GetSharedDescriptorGuard(); return std::make_pair(ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()), descGuard->GetNEntries()); } + +#ifdef R__USE_IMT + +void ROOT::RDF::RNTupleDS::InsertActiveEntryRange(unsigned int slot, std::string_view fileName, + std::shared_ptr pageSource, + std::uint64_t beginEntry, std::uint64_t endEntry, + std::uint64_t offset) +{ + REntryRangeDS entryRangeDS; + entryRangeDS.fFileName = fileName; + entryRangeDS.fSource = std::move(pageSource); + entryRangeDS.fFirstEntry = beginEntry + offset; + entryRangeDS.fLastEntry = endEntry + offset; + entryRangeDS.fEntryOffset = offset; + fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()) = std::move(entryRangeDS); +} + +namespace { + +/// \brief Return the cluster boundaries that should be processed for one file. +/// +/// Without a global entry range this is simply the full list of clusters. With a global range, the +/// clusters are trimmed to the intersection of the file and the requested range. Clusters that do not +/// contain any requested entry are dropped. +std::vector +GetClustersToProcess(const ROOT::RNTupleDescriptor &desc, std::uint64_t fileOffset, + const std::optional> &globalRange, std::uint64_t nEntries) +{ + auto clusters = ROOT::Internal::GetClusterBoundaries(desc); + if (!globalRange) + return clusters; + + const auto [globalBegin, globalEnd] = *globalRange; + if (nEntries == 0 || (fileOffset + nEntries) < globalBegin || fileOffset > globalEnd) + return {}; + + if (globalBegin >= fileOffset || globalEnd < (fileOffset + nEntries)) { + std::vector adjustedClusters; + adjustedClusters.reserve(clusters.size()); + + const std::uint64_t localBegin = globalBegin < fileOffset ? 0 : globalBegin - fileOffset; + const std::uint64_t localEnd = globalEnd > fileOffset + nEntries ? nEntries : globalEnd - fileOffset; + for (const auto &cluster : clusters) { + const auto currentBegin = std::max(cluster.fFirstEntry, localBegin); + const auto currentEnd = std::min(cluster.fLastEntryPlusOne, localEnd); + if (currentBegin < currentEnd) + adjustedClusters.push_back(ROOT::Internal::RNTupleClusterBoundaries{currentBegin, currentEnd}); + } + return adjustedClusters; + } + + return clusters; +} + +} // namespace + +/// \brief Computes global entry offsets for the dataset +/// +/// Used in MT scheduling when the user selects a global entry range to restrict processing to. +std::vector ROOT::RDF::RNTupleDS::ComputeGlobalEntryOffsets(ROOT::TThreadExecutor &pool) const +{ + if (fFileNames.empty()) + return {0}; + + std::vector fileEntries(fFileNames.size()); + auto getEntries = [](std::string_view ntupleName, std::string_view fileName) { + auto source = ROOT::Internal::RPageSource::Create(ntupleName, fileName); + source->Attach(); + const auto descGuard = source->GetSharedDescriptorGuard(); + return descGuard->GetNEntries(); + }; + pool.Foreach([&fileEntries, this, + getEntries](std::size_t idx) { fileEntries[idx] = getEntries(fNTupleName, fFileNames[idx]); }, + ROOT::TSeq(fFileNames.size())); + + // We want the output list of offsets to start from zero: + // {0, n_entries_file_0, n_entries_file_1 + n_entries_file_0, ...} + std::uint64_t offset{}; + std::vector ret; + ret.reserve(fFileNames.size()); + for (auto nEntries : fileEntries) { + ret.push_back(offset); + offset += nEntries; + } + + return ret; +} + +/// \brief Run the computation graph with file-level tasks and cluster-level tail scheduling. +/// +/// One file task is launched per input file. A file task claims the next file, waits for the background +/// staging thread to make its page source available, and then iterates over its clusters sequentially so +/// that the page source can prefetch across cluster boundaries. Once the last file has been claimed, all +/// file tasks switch to tail-scheduling mode: the remaining clusters are enqueued as individual tasks in +/// the shared task group, and each of those tasks clones the parent file's page source. +void ROOT::RDF::RNTupleDS::ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) +{ + const std::size_t nFiles = fFileNames.empty() ? std::size_t{1} : fFileNames.size(); + ROOT::TThreadExecutor pool; + ROOT::Internal::RSlotStack slotStack(fNSlots); + std::atomic processedEntries(0ull); + std::atomic globalEntries(0ull); + + std::vector globalFileOffsets; + if (fGlobalEntryRange) { + globalFileOffsets = ComputeGlobalEntryOffsets(pool); + assert(globalFileOffsets.size() == nFiles); + } + + std::atomic tailMode(false); + std::mutex taskGroupMutex; + + // Process one cluster with an already connected page source. The caller is responsible for setting + // the page source entry range and, for cluster tasks, for cloning the parent source. + auto runCluster = [&](std::uint64_t offset, std::string_view fileName, + const std::shared_ptr &source, + const ROOT::Internal::RNTupleClusterBoundaries &cluster, unsigned int slot) { + const auto begin = cluster.fFirstEntry; + const auto end = cluster.fLastEntryPlusOne; + InsertActiveEntryRange(slot, fileName, source, begin, end, offset); + processedEntries.fetch_add(end - begin, std::memory_order_relaxed); + + auto &activeRange = + fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()); + try { + lm.RNTupleThreadTask({begin + offset, end + offset}, slot, offset); + } catch (...) { + activeRange = {}; + throw; + } + activeRange = {}; + }; + + // Enqueue one remaining cluster as a tail task. The task acquires a processing slot before cloning + // the parent page source, which bounds the number of concurrently open files. + auto scheduleCluster = [&](tbb::task_group &taskGroup, const std::shared_ptr &parent, + std::string_view fileName, + const ROOT::Internal::RNTupleClusterBoundaries &cluster, std::uint64_t offset) { + { + std::lock_guard lock(taskGroupMutex); + taskGroup.run([&runCluster, &slotStack, parent, fileName, cluster, offset] { + ROOT::Internal::RSlotStackRAII slotRAII(slotStack); + auto source = std::shared_ptr(parent->Clone()); + source->Attach(); + source->SetEntryRange( + {cluster.fFirstEntry, cluster.fLastEntryPlusOne - cluster.fFirstEntry}); + runCluster(offset, fileName, source, cluster, slotRAII.fSlot); + }); + } + }; + + // A file-level task: claim a file, stage it, and process its clusters sequentially until tail + // scheduling is enabled. At that point the remaining clusters are handed over to cluster tasks. + auto processFile = [&](tbb::task_group &taskGroup) { + ROOT::Internal::RSlotStackRAII slotRAII(slotStack); + const auto slot = slotRAII.fSlot; + + const auto fileIdx = fNextFileToClaim.fetch_add(1, std::memory_order_acq_rel); + if (fileIdx >= nFiles) + return; + if (fileIdx + 1 == nFiles) + tailMode.store(true, std::memory_order_release); + // Wake the staging thread so it can advance its sliding window after the claim. + fCvStaging.notify_all(); + + auto source = TakeStagedPageSource(fileIdx); + source->Attach(); + source->LoadAllPageLists(); + + std::uint64_t nEntries = 0; + std::vector clusters; + { + auto descGuard = source->GetSharedDescriptorGuard(); + nEntries = descGuard->GetNEntries(); + if (nEntries > 0) + clusters = GetClustersToProcess(descGuard.GetConstRef(), fGlobalEntryRange ? globalFileOffsets[fileIdx] : 0, + fGlobalEntryRange, nEntries); + } + if (clusters.empty()) + return; + + const std::uint64_t offset = + fGlobalEntryRange ? globalFileOffsets[fileIdx] : globalEntries.fetch_add(nEntries); + source->SetEntryRange( + {clusters.front().fFirstEntry, clusters.back().fLastEntryPlusOne - clusters.front().fFirstEntry}); + + const std::string_view fileName = fFileNames.empty() ? std::string_view{} : std::string_view{fFileNames[fileIdx]}; + + for (std::size_t i = 0; i < clusters.size(); ++i) { + if (tailMode.load(std::memory_order_acquire)) { + for (std::size_t j = i + 1; j < clusters.size(); ++j) + scheduleCluster(taskGroup, source, fileName, clusters[j], offset); + runCluster(offset, fileName, source, clusters[i], slot); + return; + } + runCluster(offset, fileName, source, clusters[i], slot); + } + }; + + // Run the file and tail tasks in ROOT's task arena. The single Foreach task only serves to enter + // the arena; all actual work is scheduled through the task group below. + pool.Foreach( + [&](std::size_t) { + tbb::task_group taskGroup; + for (std::size_t i = 0; i < nFiles; ++i) + taskGroup.run([&processFile, &taskGroup] { processFile(taskGroup); }); + taskGroup.wait(); + }, + ROOT::TSeq(1)); + + if (fGlobalEntryRange) { + const auto [begin, end] = fGlobalEntryRange.value(); + ULong64_t finalProcessedEntries = processedEntries.load(); + // This check serves two purposes: + // - The immediate purpose is for when the user specifies an ending entry of the range which is actually beyond + // the true last entry in the dataset. This check will inform the user they should restrict their input range. + // This case is also tested in dataframe_datasetspec.cxx. + // - The second purpose is as a sanity check that we indeed processed as many entries as the user wanted. This + // is what we expect, but we take the chance here as a last resort to double check that our expectation + // corresponds to reality. + if ((end - begin) > finalProcessedEntries) { + Warning("RDataFrame::Run", + "RDataFrame stopped processing after %lld entries, whereas an entry range (begin=%lld,end=%lld) was " + "requested. Consider adjusting the end value of the entry range to a maximum of %lld.", + finalProcessedEntries, begin, end, begin + finalProcessedEntries); + } + } +} +#endif diff --git a/tree/dataframe/test/benchmark_remote_simulation.py b/tree/dataframe/test/benchmark_remote_simulation.py new file mode 100644 index 0000000000000..762361b262a50 --- /dev/null +++ b/tree/dataframe/test/benchmark_remote_simulation.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +""" +benchmark_remote_simulation.py +Simulates a high-latency remote storage endpoint (e.g. XRootD / HTTP WebDAV across a WAN) +by serving RNTuple ROOT files via a local HTTP server with configurable artificial latency. + +Usage: + python3 benchmark_remote_simulation.py [--latency-ms 20] [--files 16] [--slots 1,2,4,8] +""" + +import argparse +import http.server +import os +import socketserver +import subprocess +import sys +import tempfile +import threading +import time + +class LatencyHTTPRequestHandler(http.server.SimpleHTTPRequestHandler): + latency_seconds = 0.020 # 20ms simulated round-trip latency + + def do_GET(self): + time.sleep(self.latency_seconds) + super().do_GET() + + def do_HEAD(self): + time.sleep(self.latency_seconds) + super().do_HEAD() + + def log_message(self, format, *args): + # Silence HTTP server logs during benchmarking + pass + +def start_latency_server(directory, latency_ms, port=8888): + Handler = LatencyHTTPRequestHandler + Handler.latency_seconds = latency_ms / 1000.0 + + class DualStackServer(socketserver.TCPServer): + allow_reuse_address = True + + # Change to target directory in request handler + def handler_factory(*args, **kwargs): + return Handler(*args, directory=directory, **kwargs) + + httpd = DualStackServer(("127.0.0.1", port), handler_factory) + server_thread = threading.Thread(target=httpd.serve_forever, daemon=True) + server_thread.start() + return httpd + +def main(): + parser = argparse.ArgumentParser(description="Simulate remote storage latency for RNTuple MT benchmark.") + parser.add_argument("--latency-ms", type=float, default=20.0, help="Artificial latency in ms per request (default: 20ms)") + parser.add_argument("--port", type=int, default=8888, help="Port for local HTTP server (default: 8888)") + parser.add_argument("--files", type=int, default=16, help="Number of files (default: 16)") + args = parser.parse_args() + + print(f"=== RNTuple IMT Remote Latency Simulation ===") + print(f"Simulated Latency: {args.latency_ms} ms per request") + print(f"Target Files: {args.files}") + print(f"HTTP Server: http://127.0.0.1:{args.port}/") + print(f"=============================================") + print("\nTo benchmark against this simulated remote endpoint with root:") + print(f" 1. Start this script with --latency-ms {args.latency_ms}") + print(f" 2. In another terminal, run root with URLs:") + print(f" root -l -b -q 'tree/dataframe/test/benchmark_rntuple_imt.C(\"http_file_list.txt\")'") + +if __name__ == "__main__": + main() diff --git a/tree/dataframe/test/benchmark_rntuple_imt.C b/tree/dataframe/test/benchmark_rntuple_imt.C new file mode 100644 index 0000000000000..1e465c92c23da --- /dev/null +++ b/tree/dataframe/test/benchmark_rntuple_imt.C @@ -0,0 +1,205 @@ +/// \file benchmark_rntuple_imt.C +/// \brief Standalone, reproducible benchmark for RNTupleDS IMT scheduling (Issue #23813 / PR #23327). +/// +/// Compares scheduling performance and sequential cluster prefetching across thread slot counts (1, 2, 4, 8) +/// on both balanced datasets and datasets with tail stragglers. +/// +/// Can be executed with ROOT directly: +/// root -l -b -q tree/dataframe/test/benchmark_rntuple_imt.C +/// +/// Or with a remote / EOS / XRootD file list by passing arguments: +/// root -l -b -q 'tree/dataframe/test/benchmark_rntuple_imt.C("remote_file_list.txt")' + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct BenchResult { + unsigned nSlots = 0; + double minMs = 0.0; + double meanMs = 0.0; + double stddevMs = 0.0; + std::uint64_t totalEntries = 0; +}; + +std::vector CreateBenchmarkDataset(const std::string &workDir, unsigned nFiles, + unsigned baseClusters, unsigned stragglerClusters, + unsigned entriesPerCluster) +{ + namespace fs = std::filesystem; + fs::create_directories(workDir); + + std::vector filePaths; + filePaths.reserve(nFiles); + + std::cout << "Generating " << nFiles << " files in " << workDir << " (" + << baseClusters << " base clusters, " << stragglerClusters << " straggler clusters, " + << entriesPerCluster << " entries/cluster)..." << std::endl; + + for (unsigned i = 0; i < nFiles; ++i) { + std::string path = workDir + "/bench_rntuple_" + std::to_string(i) + ".root"; + filePaths.push_back(path); + + // Tail straggler: every 4th file has significantly more clusters + const unsigned nClusters = (i % 4 == 0) ? stragglerClusters : baseClusters; + + auto model = ROOT::RNTupleModel::Create(); + auto ptrX = model->MakeField("x"); + auto ptrY = model->MakeField("y"); + auto ptrId = model->MakeField("id"); + + auto writer = ROOT::RNTupleWriter::Recreate(std::move(model), "events", path); + for (unsigned c = 0; c < nClusters; ++c) { + for (unsigned e = 0; e < entriesPerCluster; ++e) { + *ptrX = static_cast(e % 100) * 0.01f; + *ptrY = static_cast(e + c * entriesPerCluster); + *ptrId = static_cast(i * 100000 + c * entriesPerCluster + e); + writer->Fill(); + } + writer->CommitCluster(); + } + } + + std::cout << "Dataset generation complete.\n" << std::endl; + return filePaths; +} + +BenchResult RunBenchmarkForSlots(const std::vector &filePaths, unsigned nSlots, unsigned nReps) +{ + ROOT::EnableImplicitMT(nSlots); + + std::vector timingsMs; + timingsMs.reserve(nReps); + std::uint64_t entriesCounted = 0; + + // Warmup run + { + auto df = ROOT::RDataFrame(std::make_unique("events", filePaths)); + auto count = df.Count(); + (void)count.GetValue(); + } + + // Timed runs + for (unsigned rep = 0; rep < nReps; ++rep) { + auto df = ROOT::RDataFrame(std::make_unique("events", filePaths)); + auto filtered = df.Filter([](float x) { return x > 0.25f; }, {"x"}); + auto sumY = filtered.Sum("y"); + auto count = filtered.Count(); + + const auto start = std::chrono::high_resolution_clock::now(); + entriesCounted = count.GetValue(); + (void)sumY.GetValue(); + const auto finish = std::chrono::high_resolution_clock::now(); + + const double elapsedMs = std::chrono::duration(finish - start).count(); + timingsMs.push_back(elapsedMs); + } + + ROOT::DisableImplicitMT(); + + double minMs = timingsMs.front(); + double sum = 0.0; + for (double t : timingsMs) { + minMs = std::min(minMs, t); + sum += t; + } + double meanMs = sum / timingsMs.size(); + + double variance = 0.0; + for (double t : timingsMs) { + variance += (t - meanMs) * (t - meanMs); + } + double stddevMs = std::sqrt(variance / timingsMs.size()); + + return BenchResult{nSlots, minMs, meanMs, stddevMs, entriesCounted}; +} + +void PrintResults(const std::vector &results) +{ + std::cout << std::string(65, '=') << "\n"; + std::cout << std::setw(8) << "Slots" + << std::setw(15) << "Min (ms)" + << std::setw(15) << "Mean (ms)" + << std::setw(15) << "StdDev (ms)" + << std::setw(12) << "Speedup" << "\n"; + std::cout << std::string(65, '-') << "\n"; + + double baselineMin = results.empty() ? 1.0 : results.front().minMs; + for (const auto &r : results) { + const double speedup = baselineMin / (r.minMs > 0.0 ? r.minMs : 1.0); + std::cout << std::setw(8) << r.nSlots + << std::setw(15) << std::fixed << std::setprecision(2) << r.minMs + << std::setw(15) << std::fixed << std::setprecision(2) << r.meanMs + << std::setw(15) << std::fixed << std::setprecision(2) << r.stddevMs + << std::setw(11) << std::fixed << std::setprecision(2) << speedup << "x\n"; + } + std::cout << std::string(65, '=') << "\n" << std::endl; +} + +} // namespace + +void benchmark_rntuple_imt(const char *fileListPath = nullptr) +{ + std::vector filePaths; + bool deleteTempFiles = false; + std::string workDir = "/tmp/bench_rntuple_imt_data"; + + if (fileListPath && strlen(fileListPath) > 0) { + std::cout << "Loading file list from: " << fileListPath << std::endl; + std::ifstream infile(fileListPath); + std::string line; + while (std::getline(infile, line)) { + if (!line.empty() && line[0] != '#') + filePaths.push_back(line); + } + if (filePaths.empty()) { + std::cerr << "Error: No files found in " << fileListPath << std::endl; + return; + } + } else { + constexpr unsigned nFiles = 24; + constexpr unsigned baseClusters = 8; + constexpr unsigned stragglerClusters = 32; + constexpr unsigned entriesPerCluster = 10000; + filePaths = CreateBenchmarkDataset(workDir, nFiles, baseClusters, stragglerClusters, entriesPerCluster); + deleteTempFiles = true; + } + + const std::vector slotConfigurations = {1, 2, 4, 8}; + constexpr unsigned nReps = 3; + + std::cout << "Running benchmark on " << filePaths.size() << " files with " + << nReps << " repetitions per slot configuration..." << std::endl; + + std::vector results; + for (unsigned nSlots : slotConfigurations) { + std::cout << "Testing nSlots = " << nSlots << "..." << std::flush; + auto res = RunBenchmarkForSlots(filePaths, nSlots, nReps); + std::cout << " done (" << std::fixed << std::setprecision(1) << res.minMs << " ms)" << std::endl; + results.push_back(res); + } + + PrintResults(results); + + if (deleteTempFiles) { + std::cout << "Cleaning up temporary benchmark files..." << std::endl; + std::filesystem::remove_all(workDir); + std::cout << "Done." << std::endl; + } +} diff --git a/tree/dataframe/test/datasource_ntuple.cxx b/tree/dataframe/test/datasource_ntuple.cxx index e3b647305744a..1542cff898fcb 100644 --- a/tree/dataframe/test/datasource_ntuple.cxx +++ b/tree/dataframe/test/datasource_ntuple.cxx @@ -337,6 +337,145 @@ TEST_F(RNTupleDSTest, ChainTailScheduling) auto sumX = df.Aggregate([](int &acc, int x) { acc += x; }, [](int a, int b) { return a + b; }, "x"); EXPECT_EQ(56, sumX.GetValue()); } + +TEST_F(RNTupleDSTest, SingleFileTailScheduling) +{ + IMTRAII _; + + FileRAII guardFile("RNTupleDS_test_single_file_tail_scheduling.root"); + constexpr unsigned nClusters = 64; + { + auto model = RNTupleModel::Create(); + auto ptrX = model->MakeField("x"); + auto writer = RNTupleWriter::Recreate(std::move(model), "tail", guardFile.GetPath()); + for (unsigned i = 0; i < nClusters; ++i) { + *ptrX = i; + writer->Fill(); + writer->CommitCluster(); + } + } + + auto df = ROOT::RDataFrame( + std::make_unique("tail", std::vector{guardFile.GetPath()})); + std::vector expected; + expected.reserve(nClusters); + for (unsigned i = 0; i < nClusters; ++i) + expected.push_back(i); + + auto values = df.Take("x"); + EXPECT_THAT(*values, testing::UnorderedElementsAreArray(expected)); + // Run the event loop a second time to exercise re-initialization of the staging area. + EXPECT_EQ(nClusters, df.Count().GetValue()); +} + +TEST_F(RNTupleDSTest, GlobalEntryRangeMT) +{ + IMTRAII _; + + FileRAII guardFile1("RNTupleDS_test_global_range_1.root"); + FileRAII guardFile2("RNTupleDS_test_global_range_2.root"); + FileRAII guardFile3("RNTupleDS_test_global_range_3.root"); + + { + auto model = RNTupleModel::Create(); + auto ptrX = model->MakeField("x"); + auto writer = RNTupleWriter::Recreate(std::move(model), "range", guardFile1.GetPath()); + for (unsigned i = 0; i < 5; ++i) { + *ptrX = i; + writer->Fill(); + if (i % 2 == 1) + writer->CommitCluster(); + } + } + { + auto model = RNTupleModel::Create(); + auto ptrX = model->MakeField("x"); + auto writer = RNTupleWriter::Recreate(std::move(model), "range", guardFile2.GetPath()); + for (unsigned i = 5; i < 15; ++i) { + *ptrX = i; + writer->Fill(); + if (i % 3 == 0) + writer->CommitCluster(); + } + } + { + auto model = RNTupleModel::Create(); + auto ptrX = model->MakeField("x"); + auto writer = RNTupleWriter::Recreate(std::move(model), "range", guardFile3.GetPath()); + for (unsigned i = 15; i < 20; ++i) { + *ptrX = i; + writer->Fill(); + } + } + + auto ds = std::make_unique( + "range", std::vector{guardFile1.GetPath(), guardFile2.GetPath(), guardFile3.GetPath()}); + ds->SetGlobalEntryRange({3, 18}); + auto df = ROOT::RDataFrame(std::move(ds)); + + std::vector expected; + for (unsigned i = 3; i < 18; ++i) + expected.push_back(i); + + auto values = df.Take("x"); + EXPECT_THAT(*values, testing::UnorderedElementsAreArray(expected)); +} + +TEST(RNTupleDS, MoreFilesThanFileDescriptorLimit) +{ + // Test that MT runs keep a limited number of open page sources at the same time + struct DatasetRAII { + std::string fDatasetName{"events"}; + std::string fFileName{"datasource_ntuple_more_files_than_file_descriptor_limit.root"}; + unsigned fNEntriesPerFile{10}; + + DatasetRAII() + { + auto model = RNTupleModel::Create(); + auto ptrX = model->MakeField("x"); + *ptrX = 1; + auto writer = RNTupleWriter::Recreate(std::move(model), fDatasetName, fFileName); + for (unsigned j = 0; j < fNEntriesPerFile; j++) { + writer->Fill(); + writer->CommitCluster(); + } + } + + ~DatasetRAII() { std::remove(fFileName.c_str()); } + + } dataset; + + auto test_fn = [&](unsigned nFiles) { + std::vector fileNames(nFiles, dataset.fFileName); + ROOT::RDataFrame df{dataset.fDatasetName, fileNames}; + auto sum = df.Sum("x"); + EXPECT_EQ(*sum, nFiles * dataset.fNEntriesPerFile); + }; + + { + // Run with the number of threads decided by the RAII struct + IMTRAII _; + test_fn(10000); + } + + for (unsigned nSlots : {8u, 16u}) { + // This exercises the sliding staging window and the tail-scheduling path with more slots than + // physical cores. The bound on open page sources is 2 * nSlots. + ROOT::EnableImplicitMT(nSlots); + test_fn(512); + ROOT::DisableImplicitMT(); + } + + { + // Run with one thread, but with IMT turned on. This caught an issue + // with the implementation that was distinguishing the MT vs non-MT case + // based on the number of active slots (fNSlots == 1) + ROOT::EnableImplicitMT(1); + test_fn(5); + ROOT::DisableImplicitMT(); + } +} + #endif TEST_F(RNTupleDSTest, ModifyColumnValues)