From 20bebc0a0c6706dd0714baeaaf709d9195429438 Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Mon, 19 Oct 2026 00:11:17 +0200 Subject: [PATCH 1/2] [df] Rework RNTuple MT scheduling with splittable tasks The MT scheduling of RNTupleDS is changed. The data source adopts now a similar pattern as TTreeDS, implementing a ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) method. This allows the data source to manage the creation and scheduling of multithreaded tasks, inside of which the loop manager is invoked to apply the computation graph to the range of entries in the task. A brief summary of the scheduling approach follows: 1. One task is launched in parallel per input file 2. At the beginning of the file-based task, a page source is created and the list of clusters is computed a. If the user has not specified a global entry range, it is the actual list of clusters of the page source. This happens in ScheduleTasksForFile. b. If the user has specified a global entry range, before launching the file-based tasks the data source had computed a full list of global entry offsets across all input files. The list of clusters of the page source is then adjusted according to the position of the global entry range w.r.t. the global offsets of this file. This happens in ScheduleTasksForFileWithGlobalOffset. 3. The task goes on within ScheduleTaskForPageSource, which invokes TBB's parallel_for with two arguments, one respecting the Range and the other respecting the Body requirements. See the docs respectively at https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range and https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/par_for_body. a. The Range argument is of type RPageSourceSplittableRange, a struct representing a range of entries (N consecutive clusters) tied to a specific page source to be processed. b. The first time parallel_for is invoked, an RPageSourceSplittableRange instance is passed to it which includes the whole list of clusters for that page source. c. TBB immediately calls the Range object's SplitWork method, which will split the cluster range in roughly two equal parts, the first part is re-assigned to the same object, the second part is used to create a new RPageSourceSplittableRange. d. TBB continues to split the task recursively, until a list of all non-splittable tasks is available. e. Only then, the TBB parallel_for will schedule the various calls to the ParallelForBody function, each taking one of the cluster ranges generated before. f. Within ParallelForBody, the computation graph is executed by reading data from the page source, restricting it to the assigned cluster range, by calling RLoopManager::RNTupleThreadTask. This approach delegates to the scheduling engine (TBB in this case) the responsibility to optimise the task processing and split tasks that may be running for too long (e.g. some files may contain many more entries than others). The strategy is also documented in a markdown file for posterity. --- core/imt/CMakeLists.txt | 1 + core/imt/inc/ROOT/InternalIMTUtils.hxx | 42 ++ core/imt/inc/ROOT/TThreadExecutor.hxx | 6 + core/imt/src/TThreadExecutor.cxx | 50 ++ tree/dataframe/doc/parallel_scheduling.md | 82 ++++ tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx | 7 + tree/dataframe/inc/ROOT/RNTupleDS.hxx | 112 ++++- tree/dataframe/src/RLoopManager.cxx | 34 ++ tree/dataframe/src/RNTupleDS.cxx | 477 +++++++++++++++---- tree/dataframe/test/datasource_ntuple.cxx | 48 ++ 10 files changed, 739 insertions(+), 120 deletions(-) create mode 100644 core/imt/inc/ROOT/InternalIMTUtils.hxx create mode 100644 tree/dataframe/doc/parallel_scheduling.md 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/doc/parallel_scheduling.md b/tree/dataframe/doc/parallel_scheduling.md new file mode 100644 index 0000000000000..7b99470b37f4c --- /dev/null +++ b/tree/dataframe/doc/parallel_scheduling.md @@ -0,0 +1,82 @@ +# 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 + +[//]: # (TODO: The text is meant to document the current status of the MT strategy, knowing that in the future we may adopt a different, more optimised strategy.) + +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 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. + 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. `TThreadExecutor::Foreach` is called, spawning one task per input file of the (primary) dataset. +3. The file-based tasks specialise depending on option 2c, respectively: + a. `RNTupleDS::ScheduleTasksForFile` will schedule a task reading all entries from the file. + b. `RNTupleDS::ScheduleTasksForFileWithGlobalOffset` will schedule a task with an adjusted list of clusters. The + user-provided global entry range is matched against the global offsets computed by 2c, and a new list of + clusters is generated to restrict processing of this page source according to the match. Consequently, some + tasks may be stopped immediately if the page source boundaries do not overlap with the global entry range. +4. Both file-based task implementations described above eventually call `RNTupleDS::ScheduleTaskForPageSource`. This + method redirects to the TBB parallel_for overload that takes two arguments, one respecting the + [Range](https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range) and the other respecting the [Body](https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/par_for_body) requirements. Note that the parallel_for call is hidden by + `ROOT::Internal::IMTUtils::ParallelFor`, so that the actual TBB implementation is not leaked. + a. The Range argument is of type `RSplittableRangeOfClusterBoundaries`, a struct representing a range of entries (N + consecutive clusters) tied to a specific page source to be processed. `ScheduleTaskForPageSource` calls + parallel_for with an instance of `RSplittableRangeOfClusterBoundaries` that includes the full list of valid clusters of + the corresponding page source, as defined by step 3. + b. Immediately after, TBB calls `RSplittableRangeOfClusterBoundaries::SplitWork`, which will split the + cluster range in roughly two equal parts, recursively until there are only non-splittable tasks based on the + condition given by `RSplittableRangeOfClusterBoundaries::IsDivisible`. This is the core of the strategy: we cannot know a + priori if one of the input files will be much larger than the others, or if one thread finishes earlier than + others. For this reason, we allow tasks to split themselves, so that threads that finish a file-based task + earlier can pick up one of the smaller cluster-based tasks coming as a result of + `RSplittableRangeOfClusterBoundaries::SplitWork` of another file-based task. Also, the splitting done with this strategy + ensures all the ranges of entries are non-overlapping. `RSplittableRangeOfClusterBoundaries::IsDivisible` currently + dictates that a task can split itself only if it has three or more clusters to process. The reason for requiring + at least three clusters is that the page source will prefetch the next cluster while processing the current one. + While TBB allows to provide a proportional splitting, we decide to go for the balanced basic splitting and have + a task only split into two roughly equal parts, for lack of a motivation to do differently. + c. The Body argument calls `RNTupleDS::ParallelForBody`. This function is responsible to apply the computation graph + to a page source restricted to the given range of clusters of the `RSplittableRangeOfClusterBoundaries` argument by + calling `RLoopManager::RNTupleThreadTask`. + +One cause for concern in MT runs is the risk of having too many files open at the same time, if the data source keeps +spawning tasks that open a page source. The strategy described above also limits the number of concurrently open files +through the mechanism of the RDataFrame processing slot. Here is a description of how that works, using concepts +described above with a different perspective: + +1. In the beginning of the computation graph run, the pool spawns one task per file. The pool has `N` available threads, + which start and pick one of the file-based tasks. +2. The file-based task is a function. At the beginning of this function, a page source is created, so that the valid + list of clusters of step 3 above can be computed. +3. Within `ScheduleTaskForPageSource`, parallel_for is called with the Range and Body arguments. There is no creation + of new page sources or cloning of existing ones yet. +4. Once `ParallelForBody` starts running, it will request a slot from the `ROOT::Internal::RSlotStackRAII` + argument. This effectively locks that function to run and apply the whole computation graph on the given range before + ending. Only as many `ParallelForBody` functions can run as the number of available slots. +5. Only after a slot has been taken, then `ParallelForBody` clones the input page source and restricts the cloned page + source to the assigned cluster range of the `RSplittableRangeOfClusterBoundaries` argument, with `RPageSource::SetEntryRange`. + a. In order to save one clone, one of the `ParallelForBody` will just take the input (shared) page source without + cloning it. At the moment, this is chosen to be the function execution corresponding to a + `RSplittableRangeOfClusterBoundaries` that processes the first cluster of the page source. + +In summary, creation of a page source happens in one of the two methods `ScheduleTasksForFile` or +`ScheduleTasksForFileWithGlobalOffset`. A page source created at the beginning of a file-based task can also be cloned +within a `ParallelForBody`. There are only `N` available threads, each of which could have launched a file-based task, +then end up running a `ParallelForBody`. Thus, there can be at most `2*N` open page sources at the same time. 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..b9033fad61708 100644 --- a/tree/dataframe/inc/ROOT/RNTupleDS.hxx +++ b/tree/dataframe/inc/ROOT/RNTupleDS.hxx @@ -35,6 +35,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 +87,43 @@ 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; +class TThreadExecutor; } -namespace ROOT::Internal::RDF { -class RNTupleColumnReader; -std::vector> -GetDatasetGlobalClusterBoundaries(const ROOT::RDF::RInterface &node); -} -namespace ROOT::Internal { -class RPageSource; +namespace ROOT::Internal::IMTUtils { +class RAbstractTBBRangeImpl; } +#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 +149,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 +177,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; @@ -221,7 +256,36 @@ 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); + void ParallelForBody(std::shared_ptr pageSource, + const ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl &rangeImpl, + ROOT::Detail::RDF::RLoopManager &lm, std::string_view fileName, + ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + ULong64_t columnReaderOffset); + void ScheduleTaskForPageSource(ROOT::TThreadExecutor &pool, std::unique_ptr pageSource, + std::string_view fileName, + std::span clusterBoundaries, + std::uint64_t nEntries, ROOT::Detail::RDF::RLoopManager &lm, + ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + std::atomic &globalEntries); + std::vector ComputeGlobalEntryOffsets(ROOT::TThreadExecutor &pool) const; + void ScheduleTasksForFileWithGlobalOffset(std::size_t fileIdx, const std::vector &globalFileOffsets, + ROOT::TThreadExecutor &pool, ROOT::Detail::RDF::RLoopManager &lm, + ROOT::Internal::RSlotStack &slotStack, + std::atomic &processedEntries, + std::atomic &globalEntries); + void ScheduleTasksForFile(std::size_t fileIdx, ROOT::TThreadExecutor &pool, ROOT::Detail::RDF::RLoopManager &lm, + ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + std::atomic &globalEntries); + 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..1628b5d6b7eb3 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. * @@ -37,6 +39,12 @@ #include #include #include +#ifdef R__USE_IMT +#include +#include +#include +#include +#endif // clang-format off /** @@ -303,7 +311,9 @@ ROOT::RDF::RNTupleDS::RNTupleDS(std::unique_ptr pag { pageSource->Attach(); fPrincipalDescriptor = pageSource->GetSharedDescriptorGuard()->Clone(); - fStagingArea.emplace_back(std::move(pageSource)); + // TODO(vpadulan): file staging is disabled with the current MT strategy, but it would be good to re-enable it. + if (!ROOT::IsImplicitMTEnabled()) + fStagingArea.emplace_back(std::move(pageSource)); AddField(fPrincipalDescriptor, "", fPrincipalDescriptor.GetFieldZeroId(), std::vector()); @@ -537,70 +547,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 +559,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 +672,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 +681,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,7 +729,9 @@ void ROOT::RDF::RNTupleDS::Initialize() fSeenEntriesNoGlobalRange = 0; fNextFileIndex = 0; fIsReadyForStaging = fHasNextSources = fStagingThreadShouldTerminate = false; - fThreadStaging = std::thread(&RNTupleDS::ExecStaging, this); + // TODO(vpadulan): file staging is disabled with the current MT strategy, but it would be good to re-enable it. + if (!ROOT::IsImplicitMTEnabled()) + fThreadStaging = std::thread(&RNTupleDS::ExecStaging, this); assert(fNextRanges.empty()); if (fCurrentRanges.empty() || (fFileNames.size() > fNSlots)) { @@ -818,7 +760,9 @@ void ROOT::RDF::RNTupleDS::Finalize() fStagingThreadShouldTerminate = true; } fCvStaging.notify_one(); - fThreadStaging.join(); + // TODO(vpadulan): file staging is disabled with the current MT strategy, but it would be good to re-enable it. + if (!ROOT::IsImplicitMTEnabled()) + 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 +779,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 +800,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 +845,340 @@ 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 +namespace { + +/// \brief Implementation of the Range type to be used in a parallel_for when processing RNTuple data with RDataFrame. +/// +/// The class represents one range of entries (N consecutive clusters) tied to a specific page source to be processed. +/// The current MT strategy launches an outer parallel_for with exactly one task per page source, then within each of +/// these launches another parallel_for that will take an instance of this class. The initial state of the class is to +/// include the whole list of clusters of the corresponding page source. TBB will eagerly split this range in as many +/// ranges as possible, until there is a full set of non-splittable ranges (based on IsDivisible). The split ranges +/// only have a different begin and end cluster, there is no page source copying or cloning done by this range, instead +/// that happens in the ParallelForBody so that it can be limited by the number of RDF slots used. +class RSplittableRangeOfClusterBoundaries final : public ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl { + // A view of all the cluster boundaries in the page source + std::span fClusters; + std::size_t fClusterIdxBegin; // inclusive + std::size_t fClusterIdxEnd; // exclusive + +public: + RSplittableRangeOfClusterBoundaries(std::span clusters, + std::size_t clusterIdxBegin, std::size_t clusterIdxEnd) + : fClusters(clusters), fClusterIdxBegin(clusterIdxBegin), fClusterIdxEnd(clusterIdxEnd) + { + assert(fClusterIdxEnd > fClusterIdxBegin); + } + + std::unique_ptr SplitWork() final + { + // Split the cluster range in half + auto clusterIdxMid = fClusterIdxBegin + (fClusterIdxEnd - fClusterIdxBegin) / 2; + auto newClusterIdxBegin = clusterIdxMid; + auto newClusterIdxEnd = fClusterIdxEnd; + fClusterIdxEnd = clusterIdxMid; + + return std::make_unique(fClusters, newClusterIdxBegin, newClusterIdxEnd); + } + + bool IsDivisible() const final + { + // We choose two because the page source prefetches one cluster ahead of the current one. + return fClusterIdxEnd - fClusterIdxBegin > 2; + } + + bool Empty() const final + { + // empty is called once at the beginning of the parallel_for, to guard against empty tasks. We only pass tasks + // with actual work by construction + return false; + } + + std::uint64_t GetCurrentBeginEntry() const { return fClusters[fClusterIdxBegin].fFirstEntry; } + + std::uint64_t GetCurrentEndEntry() const { return fClusters[fClusterIdxEnd - 1].fLastEntryPlusOne; } + + bool IsProcessingFirstCluster() const { return fClusterIdxBegin == 0; } +}; +} // namespace + +// For now we need this function as a pass-through for the information contained in a +// RSplittableRangeOfClusterBoundaries to be available to the rest of the computation graph (e.g. to connect the column +// readers and update the sample info) +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); +} + +/// \brief Apply the computation graph to an input range. +/// \param pageSource The page source to read data from. +/// \param rangeImpl A Range type carrying information on the range of entries to process. +/// \param lm The RLoopManager which will run the computation graph. +/// \param fileName Path to the storage location of the page source. +/// \param slotStack A stack of slots, each thread needs to request one before running the computation graph. +/// \param processedEntries Cumulating number of processed entries by this computation graph run. +/// \param columnReaderOffset Offset value to pass to column readers used in this task. The column reader will request +/// entries from the page source receiving a global entry number, this offset is subtracted so that the column reader +/// requests the local entry from the page source. +void ROOT::RDF::RNTupleDS::ParallelForBody(std::shared_ptr pageSource, + const ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl &rangeImpl, + ROOT::Detail::RDF::RLoopManager &lm, std::string_view fileName, + ROOT::Internal::RSlotStack &slotStack, + std::atomic &processedEntries, ULong64_t columnReaderOffset) +{ + ROOT::Internal::RSlotStackRAII slotRAII{slotStack}; + const auto &slot = slotRAII.fSlot; + + const auto &range = static_cast(rangeImpl); + const auto &beginEntry = range.GetCurrentBeginEntry(); + const auto &endEntry = range.GetCurrentEndEntry(); + + // We clone the page source only within this function scope, to limit the number of open page sources to 2*fNSlots + auto localPageSource = [&]() -> std::shared_ptr { + // In order to avoid an extra Clone, we need to choose one of the ranges that will take the original page source, + // we choose it to be the one processing the first valid cluster + if (range.IsProcessingFirstCluster()) + return pageSource; + + auto newPageSource = pageSource->Clone(); + newPageSource->Attach(); + return newPageSource; + }(); + + // Notify the current page source about the current entry range + localPageSource->SetEntryRange({beginEntry, endEntry - beginEntry}); + + processedEntries.fetch_add(endEntry - beginEntry); + + // At this point we let the RNTupleDS know about the range being processed by this task. This is used for relevant + // operations e.g. connecting the column readers as well as populating sample information. + InsertActiveEntryRange(slot, fileName, std::move(localPageSource), beginEntry, endEntry, columnReaderOffset); + + lm.RNTupleThreadTask({beginEntry + columnReaderOffset, endEntry + columnReaderOffset}, slot, columnReaderOffset); + + // Clear the active range after having processed it, to avoid keeping around the page source for longer than intended + fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()) = {}; +} + +/// \brief Call parallel_for to schedule a splittable task for a page source +/// \param pool The pool of threads used to schedule the task. +/// \param pageSource The RNTuple page source for the task. +/// \param fileName Path to the storage location of the page source. +/// \param clusterBoundaries A list of cluster boundaries valid for the page source. This may be computed differently, +/// see ProcessMT. +/// \param nEntries Number of entries in the page source. +/// \param lm The RLoopManager which will run the computation graph. +/// \param slotStack A stack of slots, each thread needs to request one before running the computation graph. +/// \param processedEntries Cumulating number of processed entries by this computation graph run. +/// \param globalEntries Cumulating total number of entries in the input dataset. +/// +/// This function calls the parallel_for overload that takes two arguments, one respecting the Range and the other +/// respecting the Body requirements. The Range argument is a RSplittableRangeOfClusterBoundaries, a struct with a range +/// of N cluster boundaries to process (initialised to the whole list of clusters). Immediately when calling +/// parallel_for, TBB will recursively ask the RSplittableRangeOfClusterBoundaries to split itself into two roughly +/// equal sub-tasks, until there are only leaf tasks that cannot be split further. This helps balancing cases where some +/// tasks may end up running much longer than others. +void ROOT::RDF::RNTupleDS::ScheduleTaskForPageSource( + ROOT::TThreadExecutor &pool, std::unique_ptr pageSource, std::string_view fileName, + std::span clusterBoundaries, std::uint64_t nEntries, + ROOT::Detail::RDF::RLoopManager &lm, ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + std::atomic &globalEntries) +{ + + const auto columnReaderOffset{globalEntries.fetch_add(nEntries)}; + + // We need to share the same page source pointer across all splitted tasks, to decide if we clone it (in most cases) + // or not (one task will take the original page source). + std::shared_ptr sharedPageSource{std::move(pageSource)}; + + auto parallelForBody = [this, &sharedPageSource, &lm, &fileName, &slotStack, &processedEntries, + &columnReaderOffset](const ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl &rangeImpl) { + this->ParallelForBody(sharedPageSource, rangeImpl, lm, fileName, slotStack, processedEntries, columnReaderOffset); + }; + + ROOT::Internal::IMTUtils::ParallelFor( + pool, std::make_unique(clusterBoundaries, 0, clusterBoundaries.size()), + parallelForBody); +} + +/// \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 +{ + 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 Launch the task(s) for the input file index, taking into account dataset global entry offsets. +/// +/// This function is used in MT scheduling to launch a task processing one input file. It is selected when +/// the user specifies a global entry range to restrict processing to. In this scenario we have precomputed the global +/// entry offsets of each file upfront so we can then adjust the range of entries in each file accordingly. The task may +/// be split into smaller tasks based on entry ranges aligned to cluster boundaries, see ScheduleTaskForPageSource. +void ROOT::RDF::RNTupleDS::ScheduleTasksForFileWithGlobalOffset( + std::size_t fileIdx, const std::vector &globalFileOffsets, ROOT::TThreadExecutor &pool, + ROOT::Detail::RDF::RLoopManager &lm, ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + std::atomic &globalEntries) +{ + const auto &fileName = fFileNames[fileIdx]; + auto source = CreatePageSource(fNTupleName, fileName); + source->Attach(); + // TODO: improve MT scheduling such that not all page lists are required right away + source->LoadAllPageLists(); + + const std::uint64_t globalFileOffset = globalFileOffsets[fileIdx]; + const auto [globalBegin, globalEnd] = fGlobalEntryRange.value(); + + const auto clustersAndEntries = + [&]() -> std::pair, ROOT::NTupleSize_t> { + const auto descGuard = source->GetSharedDescriptorGuard(); + const auto nEntries = descGuard->GetNEntries(); + // Early return: + // - There are no entries in the file + // - The boundaries of this file are not touched by the global entry range + if (nEntries == 0 || (globalFileOffset + nEntries) < globalBegin || globalFileOffset > globalEnd) + return {{}, ROOT::kInvalidNTupleIndex}; + + const auto clustersInFile = ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()); + + // The list of cluster boundaries in this file needs to be adjusted if either the beginning or the end of the + // global entry range fall within the boundaries of the entries of this file. + if ((globalBegin >= globalFileOffset) || (globalEnd < (globalFileOffset + nEntries))) { + std::vector adjustedClusters; + for (const auto &cluster : clustersInFile) { + const std::uint64_t localBegin = globalBegin < globalFileOffset ? 0 : globalBegin - globalFileOffset; + const std::uint64_t localEnd = + globalEnd > globalFileOffset + nEntries ? globalFileOffset + nEntries : globalEnd - globalFileOffset; + const auto currentBegin = std::max(cluster.fFirstEntry, localBegin); + const auto currentEnd = std::min(cluster.fLastEntryPlusOne, localEnd); + // This is not satified if the desired begin is larger than the last entry of some cluster + // In this case, this cluster is not going to be processed further + if (currentBegin < currentEnd) + adjustedClusters.push_back(ROOT::Internal::RNTupleClusterBoundaries{currentBegin, currentEnd}); + } + return std::make_pair(adjustedClusters, nEntries); + } + + return std::make_pair(clustersInFile, nEntries); + }(); + const ROOT::NTupleSize_t nEntries = clustersAndEntries.second; + if (nEntries == ROOT::kInvalidNTupleIndex) + return; + const auto &clusters = clustersAndEntries.first; + ScheduleTaskForPageSource(pool, std::move(source), fileName, clusters, nEntries, lm, slotStack, processedEntries, + globalEntries); +} + +/// \brief Launch the task(s) for the input file index +/// +/// This function is used in MT scheduling to launch a task processing one input file. It is selected when +/// there is no global entry range. The task may be split into smaller tasks based on entry ranges aligned to cluster +/// boundaries, see ScheduleTaskForPageSource. +void ROOT::RDF::RNTupleDS::ScheduleTasksForFile(std::size_t fileIdx, ROOT::TThreadExecutor &pool, + ROOT::Detail::RDF::RLoopManager &lm, + ROOT::Internal::RSlotStack &slotStack, + std::atomic &processedEntries, + std::atomic &globalEntries) +{ + // Evaluate clusters (with local entry numbers) and number of entries for this file + const auto &fileName = fFileNames[fileIdx]; + auto source = CreatePageSource(fNTupleName, fFileNames[fileIdx]); + source->Attach(); + // TODO: improve MT scheduling such that not all page lists are required right away + source->LoadAllPageLists(); + const auto clustersAndEntries = [&]() { + const auto descGuard = source->GetSharedDescriptorGuard(); + return std::make_pair(ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()), descGuard->GetNEntries()); + }(); + const ROOT::NTupleSize_t nEntries = clustersAndEntries.second; + if (nEntries == 0) + return; + const auto &clusters = clustersAndEntries.first; + ScheduleTaskForPageSource(pool, std::move(source), fileName, clusters, nEntries, lm, slotStack, processedEntries, + globalEntries); +} + +/// \brief Run the computation graph by scheduling MT tasks +/// +/// The data source interacts with the loop manager and with a pool of threads (via TThreadExecutor) at the same time. +/// The function schedules one MT task per input file, eventually calling the parallel_for overload taking two input +/// arguments respectively respecting the Range and Body requirements. This means that each of the file-based tasks +/// can be split into various sub-tasks, each processing a set of consecutive clusters of the same file. Before +/// launching the file-based tasks, the function checks if the user has selected a global range of entries to restrict +/// processing to. If so, the function first computes the global entry offsets of all the input files, so they can be +/// matched against the user-provided global range. These will be used to adjust the list of clusters of the file before +/// spawning the corresponding task. +void ROOT::RDF::RNTupleDS::ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) +{ + ROOT::TThreadExecutor pool; + ROOT::Internal::RSlotStack slotStack(fNSlots); + std::atomic processedEntries(0ull); + std::atomic globalEntries(0ull); + + if (fGlobalEntryRange.has_value()) { + const auto globalFileOffsets = ComputeGlobalEntryOffsets(pool); + pool.Foreach( + [&](std::size_t fileIdx) { + ScheduleTasksForFileWithGlobalOffset(fileIdx, globalFileOffsets, pool, lm, slotStack, processedEntries, + globalEntries); + }, + ROOT::TSeq(fFileNames.size())); + } else { + pool.Foreach( + [&](std::size_t fileIdx) { + ScheduleTasksForFile(fileIdx, pool, lm, slotStack, processedEntries, globalEntries); + }, + ROOT::TSeq(fFileNames.size())); + } + + if (fGlobalEntryRange.has_value()) { + 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/datasource_ntuple.cxx b/tree/dataframe/test/datasource_ntuple.cxx index e3b647305744a..7645f7026b85c 100644 --- a/tree/dataframe/test/datasource_ntuple.cxx +++ b/tree/dataframe/test/datasource_ntuple.cxx @@ -337,6 +337,54 @@ 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(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); + } + + { + // 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) From 1c3d0644fe578666300d22a620db251ffb4f79cd Mon Sep 17 00:00:00 2001 From: Golam Samdani Sayam Date: Fri, 23 Oct 2026 22:46:58 +0600 Subject: [PATCH 2/2] [df] Improve RNTuple IMT scheduling with sliding-window staging and tail splitting Follow-up to PR #23327 (Issue #23813). Replaces TBB splittable ranges with file-level task distribution, background sliding-window staging, and dynamic tail-cluster scheduling: - One file-level task per input file dispatched into a tbb::task_group. - Sequential cluster iteration per file preserves read-ahead and cluster prefetching across cluster boundaries. - Staging thread prefetches page sources ahead of the next file to be claimed with a window bounded by fNSlots. - Transition into tail-cluster scheduling once all files have been claimed; remaining clusters are scheduled individually across idle workers. - Strictly bounds concurrent open page sources to at most 2 * nSlots. - Added tests for single-file tail scheduling, global entry ranges, and file descriptor bounds with high slot counts (8/16 slots, 512 files). - Documented new architecture in tree/dataframe/doc/parallel_scheduling.md. - Added reproducible benchmark suite (benchmark_rntuple_imt.C and benchmark_remote_simulation.py). Resolves root-project/root#23813. --- tree/dataframe/CMakeLists.txt | 4 + tree/dataframe/doc/parallel_scheduling.md | 101 ++-- tree/dataframe/inc/ROOT/RNTupleDS.hxx | 41 +- tree/dataframe/src/RNTupleDS.cxx | 525 +++++++++--------- .../test/benchmark_remote_simulation.py | 70 +++ tree/dataframe/test/benchmark_rntuple_imt.C | 205 +++++++ tree/dataframe/test/datasource_ntuple.cxx | 91 +++ 7 files changed, 689 insertions(+), 348 deletions(-) create mode 100644 tree/dataframe/test/benchmark_remote_simulation.py create mode 100644 tree/dataframe/test/benchmark_rntuple_imt.C 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 index 7b99470b37f4c..a73f4d4acb1be 100644 --- a/tree/dataframe/doc/parallel_scheduling.md +++ b/tree/dataframe/doc/parallel_scheduling.md @@ -9,74 +9,59 @@ RDataFrame allows to process input data in parallel in two ways: with implicit m ### RNTuple -[//]: # (TODO: The text is meant to document the current status of the MT strategy, knowing that in the future we may adopt a different, more optimised strategy.) - 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 +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 methods described from now on, are private to the `RNTupleDS` class. + 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. + 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. `TThreadExecutor::Foreach` is called, spawning one task per input file of the (primary) dataset. -3. The file-based tasks specialise depending on option 2c, respectively: - a. `RNTupleDS::ScheduleTasksForFile` will schedule a task reading all entries from the file. - b. `RNTupleDS::ScheduleTasksForFileWithGlobalOffset` will schedule a task with an adjusted list of clusters. The - user-provided global entry range is matched against the global offsets computed by 2c, and a new list of - clusters is generated to restrict processing of this page source according to the match. Consequently, some - tasks may be stopped immediately if the page source boundaries do not overlap with the global entry range. -4. Both file-based task implementations described above eventually call `RNTupleDS::ScheduleTaskForPageSource`. This - method redirects to the TBB parallel_for overload that takes two arguments, one respecting the - [Range](https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range) and the other respecting the [Body](https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/par_for_body) requirements. Note that the parallel_for call is hidden by - `ROOT::Internal::IMTUtils::ParallelFor`, so that the actual TBB implementation is not leaked. - a. The Range argument is of type `RSplittableRangeOfClusterBoundaries`, a struct representing a range of entries (N - consecutive clusters) tied to a specific page source to be processed. `ScheduleTaskForPageSource` calls - parallel_for with an instance of `RSplittableRangeOfClusterBoundaries` that includes the full list of valid clusters of - the corresponding page source, as defined by step 3. - b. Immediately after, TBB calls `RSplittableRangeOfClusterBoundaries::SplitWork`, which will split the - cluster range in roughly two equal parts, recursively until there are only non-splittable tasks based on the - condition given by `RSplittableRangeOfClusterBoundaries::IsDivisible`. This is the core of the strategy: we cannot know a - priori if one of the input files will be much larger than the others, or if one thread finishes earlier than - others. For this reason, we allow tasks to split themselves, so that threads that finish a file-based task - earlier can pick up one of the smaller cluster-based tasks coming as a result of - `RSplittableRangeOfClusterBoundaries::SplitWork` of another file-based task. Also, the splitting done with this strategy - ensures all the ranges of entries are non-overlapping. `RSplittableRangeOfClusterBoundaries::IsDivisible` currently - dictates that a task can split itself only if it has three or more clusters to process. The reason for requiring - at least three clusters is that the page source will prefetch the next cluster while processing the current one. - While TBB allows to provide a proportional splitting, we decide to go for the balanced basic splitting and have - a task only split into two roughly equal parts, for lack of a motivation to do differently. - c. The Body argument calls `RNTupleDS::ParallelForBody`. This function is responsible to apply the computation graph - to a page source restricted to the given range of clusters of the `RSplittableRangeOfClusterBoundaries` argument by - calling `RLoopManager::RNTupleThreadTask`. - -One cause for concern in MT runs is the risk of having too many files open at the same time, if the data source keeps -spawning tasks that open a page source. The strategy described above also limits the number of concurrently open files -through the mechanism of the RDataFrame processing slot. Here is a description of how that works, using concepts -described above with a different perspective: + 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. -1. In the beginning of the computation graph run, the pool spawns one task per file. The pool has `N` available threads, - which start and pick one of the file-based tasks. -2. The file-based task is a function. At the beginning of this function, a page source is created, so that the valid - list of clusters of step 3 above can be computed. -3. Within `ScheduleTaskForPageSource`, parallel_for is called with the Range and Body arguments. There is no creation - of new page sources or cloning of existing ones yet. -4. Once `ParallelForBody` starts running, it will request a slot from the `ROOT::Internal::RSlotStackRAII` - argument. This effectively locks that function to run and apply the whole computation graph on the given range before - ending. Only as many `ParallelForBody` functions can run as the number of available slots. -5. Only after a slot has been taken, then `ParallelForBody` clones the input page source and restricts the cloned page - source to the assigned cluster range of the `RSplittableRangeOfClusterBoundaries` argument, with `RPageSource::SetEntryRange`. - a. In order to save one clone, one of the `ParallelForBody` will just take the input (shared) page source without - cloning it. At the moment, this is chosen to be the function execution corresponding to a - `RSplittableRangeOfClusterBoundaries` that processes the first cluster of the page source. +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`. -In summary, creation of a page source happens in one of the two methods `ScheduleTasksForFile` or -`ScheduleTasksForFileWithGlobalOffset`. A page source created at the beginning of a file-based task can also be cloned -within a `ParallelForBody`. There are only `N` available threads, each of which could have launched a file-based task, -then end up running a `ParallelForBody`. Thus, there can be at most `2*N` open page sources at the same time. +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/RNTupleDS.hxx b/tree/dataframe/inc/ROOT/RNTupleDS.hxx index b9033fad61708..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 @@ -91,9 +93,6 @@ GetClustersAndEntries(std::string_view ntupleName, std::string_view location); namespace ROOT { class TThreadExecutor; } -namespace ROOT::Internal::IMTUtils { -class RAbstractTBBRangeImpl; -} #endif namespace ROOT::RDF { @@ -199,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 { @@ -236,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 @@ -263,26 +279,7 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { 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); - void ParallelForBody(std::shared_ptr pageSource, - const ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl &rangeImpl, - ROOT::Detail::RDF::RLoopManager &lm, std::string_view fileName, - ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, - ULong64_t columnReaderOffset); - void ScheduleTaskForPageSource(ROOT::TThreadExecutor &pool, std::unique_ptr pageSource, - std::string_view fileName, - std::span clusterBoundaries, - std::uint64_t nEntries, ROOT::Detail::RDF::RLoopManager &lm, - ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, - std::atomic &globalEntries); std::vector ComputeGlobalEntryOffsets(ROOT::TThreadExecutor &pool) const; - void ScheduleTasksForFileWithGlobalOffset(std::size_t fileIdx, const std::vector &globalFileOffsets, - ROOT::TThreadExecutor &pool, ROOT::Detail::RDF::RLoopManager &lm, - ROOT::Internal::RSlotStack &slotStack, - std::atomic &processedEntries, - std::atomic &globalEntries); - void ScheduleTasksForFile(std::size_t fileIdx, ROOT::TThreadExecutor &pool, ROOT::Detail::RDF::RLoopManager &lm, - ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, - std::atomic &globalEntries); void ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) final; #endif diff --git a/tree/dataframe/src/RNTupleDS.cxx b/tree/dataframe/src/RNTupleDS.cxx index 1628b5d6b7eb3..4c43ae1fed8bb 100644 --- a/tree/dataframe/src/RNTupleDS.cxx +++ b/tree/dataframe/src/RNTupleDS.cxx @@ -29,6 +29,7 @@ #include #include +#include #include #include #include @@ -40,10 +41,9 @@ #include #include #ifdef R__USE_IMT -#include #include -#include #include +#include #endif // clang-format off @@ -311,9 +311,9 @@ ROOT::RDF::RNTupleDS::RNTupleDS(std::unique_ptr pag { pageSource->Attach(); fPrincipalDescriptor = pageSource->GetSharedDescriptorGuard()->Clone(); - // TODO(vpadulan): file staging is disabled with the current MT strategy, but it would be good to re-enable it. - if (!ROOT::IsImplicitMTEnabled()) - fStagingArea.emplace_back(std::move(pageSource)); + // 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(), std::vector()); @@ -512,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()); @@ -729,11 +789,27 @@ void ROOT::RDF::RNTupleDS::Initialize() fSeenEntriesNoGlobalRange = 0; fNextFileIndex = 0; fIsReadyForStaging = fHasNextSources = fStagingThreadShouldTerminate = false; - // TODO(vpadulan): file staging is disabled with the current MT strategy, but it would be good to re-enable it. - if (!ROOT::IsImplicitMTEnabled()) - 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. { @@ -759,9 +835,30 @@ void ROOT::RDF::RNTupleDS::Finalize() std::lock_guard _(fMutexStaging); fStagingThreadShouldTerminate = true; } - fCvStaging.notify_one(); - // TODO(vpadulan): file staging is disabled with the current MT strategy, but it would be good to re-enable it. - if (!ROOT::IsImplicitMTEnabled()) + 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. @@ -847,66 +944,7 @@ ROOT::Internal::RDF::GetClustersAndEntries(std::string_view ntupleName, std::str } #ifdef R__USE_IMT -namespace { - -/// \brief Implementation of the Range type to be used in a parallel_for when processing RNTuple data with RDataFrame. -/// -/// The class represents one range of entries (N consecutive clusters) tied to a specific page source to be processed. -/// The current MT strategy launches an outer parallel_for with exactly one task per page source, then within each of -/// these launches another parallel_for that will take an instance of this class. The initial state of the class is to -/// include the whole list of clusters of the corresponding page source. TBB will eagerly split this range in as many -/// ranges as possible, until there is a full set of non-splittable ranges (based on IsDivisible). The split ranges -/// only have a different begin and end cluster, there is no page source copying or cloning done by this range, instead -/// that happens in the ParallelForBody so that it can be limited by the number of RDF slots used. -class RSplittableRangeOfClusterBoundaries final : public ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl { - // A view of all the cluster boundaries in the page source - std::span fClusters; - std::size_t fClusterIdxBegin; // inclusive - std::size_t fClusterIdxEnd; // exclusive - -public: - RSplittableRangeOfClusterBoundaries(std::span clusters, - std::size_t clusterIdxBegin, std::size_t clusterIdxEnd) - : fClusters(clusters), fClusterIdxBegin(clusterIdxBegin), fClusterIdxEnd(clusterIdxEnd) - { - assert(fClusterIdxEnd > fClusterIdxBegin); - } - - std::unique_ptr SplitWork() final - { - // Split the cluster range in half - auto clusterIdxMid = fClusterIdxBegin + (fClusterIdxEnd - fClusterIdxBegin) / 2; - auto newClusterIdxBegin = clusterIdxMid; - auto newClusterIdxEnd = fClusterIdxEnd; - fClusterIdxEnd = clusterIdxMid; - - return std::make_unique(fClusters, newClusterIdxBegin, newClusterIdxEnd); - } - - bool IsDivisible() const final - { - // We choose two because the page source prefetches one cluster ahead of the current one. - return fClusterIdxEnd - fClusterIdxBegin > 2; - } - - bool Empty() const final - { - // empty is called once at the beginning of the parallel_for, to guard against empty tasks. We only pass tasks - // with actual work by construction - return false; - } - - std::uint64_t GetCurrentBeginEntry() const { return fClusters[fClusterIdxBegin].fFirstEntry; } - - std::uint64_t GetCurrentEndEntry() const { return fClusters[fClusterIdxEnd - 1].fLastEntryPlusOne; } - - bool IsProcessingFirstCluster() const { return fClusterIdxBegin == 0; } -}; -} // namespace -// For now we need this function as a pass-through for the information contained in a -// RSplittableRangeOfClusterBoundaries to be available to the rest of the computation graph (e.g. to connect the column -// readers and update the sample info) void ROOT::RDF::RNTupleDS::InsertActiveEntryRange(unsigned int slot, std::string_view fileName, std::shared_ptr pageSource, std::uint64_t beginEntry, std::uint64_t endEntry, @@ -921,102 +959,53 @@ void ROOT::RDF::RNTupleDS::InsertActiveEntryRange(unsigned int slot, std::string fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()) = std::move(entryRangeDS); } -/// \brief Apply the computation graph to an input range. -/// \param pageSource The page source to read data from. -/// \param rangeImpl A Range type carrying information on the range of entries to process. -/// \param lm The RLoopManager which will run the computation graph. -/// \param fileName Path to the storage location of the page source. -/// \param slotStack A stack of slots, each thread needs to request one before running the computation graph. -/// \param processedEntries Cumulating number of processed entries by this computation graph run. -/// \param columnReaderOffset Offset value to pass to column readers used in this task. The column reader will request -/// entries from the page source receiving a global entry number, this offset is subtracted so that the column reader -/// requests the local entry from the page source. -void ROOT::RDF::RNTupleDS::ParallelForBody(std::shared_ptr pageSource, - const ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl &rangeImpl, - ROOT::Detail::RDF::RLoopManager &lm, std::string_view fileName, - ROOT::Internal::RSlotStack &slotStack, - std::atomic &processedEntries, ULong64_t columnReaderOffset) -{ - ROOT::Internal::RSlotStackRAII slotRAII{slotStack}; - const auto &slot = slotRAII.fSlot; - - const auto &range = static_cast(rangeImpl); - const auto &beginEntry = range.GetCurrentBeginEntry(); - const auto &endEntry = range.GetCurrentEndEntry(); - - // We clone the page source only within this function scope, to limit the number of open page sources to 2*fNSlots - auto localPageSource = [&]() -> std::shared_ptr { - // In order to avoid an extra Clone, we need to choose one of the ranges that will take the original page source, - // we choose it to be the one processing the first valid cluster - if (range.IsProcessingFirstCluster()) - return pageSource; - - auto newPageSource = pageSource->Clone(); - newPageSource->Attach(); - return newPageSource; - }(); - - // Notify the current page source about the current entry range - localPageSource->SetEntryRange({beginEntry, endEntry - beginEntry}); - - processedEntries.fetch_add(endEntry - beginEntry); - - // At this point we let the RNTupleDS know about the range being processed by this task. This is used for relevant - // operations e.g. connecting the column readers as well as populating sample information. - InsertActiveEntryRange(slot, fileName, std::move(localPageSource), beginEntry, endEntry, columnReaderOffset); - - lm.RNTupleThreadTask({beginEntry + columnReaderOffset, endEntry + columnReaderOffset}, slot, columnReaderOffset); - - // Clear the active range after having processed it, to avoid keeping around the page source for longer than intended - fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()) = {}; -} +namespace { -/// \brief Call parallel_for to schedule a splittable task for a page source -/// \param pool The pool of threads used to schedule the task. -/// \param pageSource The RNTuple page source for the task. -/// \param fileName Path to the storage location of the page source. -/// \param clusterBoundaries A list of cluster boundaries valid for the page source. This may be computed differently, -/// see ProcessMT. -/// \param nEntries Number of entries in the page source. -/// \param lm The RLoopManager which will run the computation graph. -/// \param slotStack A stack of slots, each thread needs to request one before running the computation graph. -/// \param processedEntries Cumulating number of processed entries by this computation graph run. -/// \param globalEntries Cumulating total number of entries in the input dataset. +/// \brief Return the cluster boundaries that should be processed for one file. /// -/// This function calls the parallel_for overload that takes two arguments, one respecting the Range and the other -/// respecting the Body requirements. The Range argument is a RSplittableRangeOfClusterBoundaries, a struct with a range -/// of N cluster boundaries to process (initialised to the whole list of clusters). Immediately when calling -/// parallel_for, TBB will recursively ask the RSplittableRangeOfClusterBoundaries to split itself into two roughly -/// equal sub-tasks, until there are only leaf tasks that cannot be split further. This helps balancing cases where some -/// tasks may end up running much longer than others. -void ROOT::RDF::RNTupleDS::ScheduleTaskForPageSource( - ROOT::TThreadExecutor &pool, std::unique_ptr pageSource, std::string_view fileName, - std::span clusterBoundaries, std::uint64_t nEntries, - ROOT::Detail::RDF::RLoopManager &lm, ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, - std::atomic &globalEntries) +/// 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; + } - const auto columnReaderOffset{globalEntries.fetch_add(nEntries)}; - - // We need to share the same page source pointer across all splitted tasks, to decide if we clone it (in most cases) - // or not (one task will take the original page source). - std::shared_ptr sharedPageSource{std::move(pageSource)}; - - auto parallelForBody = [this, &sharedPageSource, &lm, &fileName, &slotStack, &processedEntries, - &columnReaderOffset](const ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl &rangeImpl) { - this->ParallelForBody(sharedPageSource, rangeImpl, lm, fileName, slotStack, processedEntries, columnReaderOffset); - }; - - ROOT::Internal::IMTUtils::ParallelFor( - pool, std::make_unique(clusterBoundaries, 0, clusterBoundaries.size()), - parallelForBody); + 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); @@ -1041,129 +1030,129 @@ std::vector ROOT::RDF::RNTupleDS::ComputeGlobalEntryOffsets(ROOT: return ret; } -/// \brief Launch the task(s) for the input file index, taking into account dataset global entry offsets. -/// -/// This function is used in MT scheduling to launch a task processing one input file. It is selected when -/// the user specifies a global entry range to restrict processing to. In this scenario we have precomputed the global -/// entry offsets of each file upfront so we can then adjust the range of entries in each file accordingly. The task may -/// be split into smaller tasks based on entry ranges aligned to cluster boundaries, see ScheduleTaskForPageSource. -void ROOT::RDF::RNTupleDS::ScheduleTasksForFileWithGlobalOffset( - std::size_t fileIdx, const std::vector &globalFileOffsets, ROOT::TThreadExecutor &pool, - ROOT::Detail::RDF::RLoopManager &lm, ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, - std::atomic &globalEntries) -{ - const auto &fileName = fFileNames[fileIdx]; - auto source = CreatePageSource(fNTupleName, fileName); - source->Attach(); - // TODO: improve MT scheduling such that not all page lists are required right away - source->LoadAllPageLists(); - - const std::uint64_t globalFileOffset = globalFileOffsets[fileIdx]; - const auto [globalBegin, globalEnd] = fGlobalEntryRange.value(); - - const auto clustersAndEntries = - [&]() -> std::pair, ROOT::NTupleSize_t> { - const auto descGuard = source->GetSharedDescriptorGuard(); - const auto nEntries = descGuard->GetNEntries(); - // Early return: - // - There are no entries in the file - // - The boundaries of this file are not touched by the global entry range - if (nEntries == 0 || (globalFileOffset + nEntries) < globalBegin || globalFileOffset > globalEnd) - return {{}, ROOT::kInvalidNTupleIndex}; - - const auto clustersInFile = ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()); - - // The list of cluster boundaries in this file needs to be adjusted if either the beginning or the end of the - // global entry range fall within the boundaries of the entries of this file. - if ((globalBegin >= globalFileOffset) || (globalEnd < (globalFileOffset + nEntries))) { - std::vector adjustedClusters; - for (const auto &cluster : clustersInFile) { - const std::uint64_t localBegin = globalBegin < globalFileOffset ? 0 : globalBegin - globalFileOffset; - const std::uint64_t localEnd = - globalEnd > globalFileOffset + nEntries ? globalFileOffset + nEntries : globalEnd - globalFileOffset; - const auto currentBegin = std::max(cluster.fFirstEntry, localBegin); - const auto currentEnd = std::min(cluster.fLastEntryPlusOne, localEnd); - // This is not satified if the desired begin is larger than the last entry of some cluster - // In this case, this cluster is not going to be processed further - if (currentBegin < currentEnd) - adjustedClusters.push_back(ROOT::Internal::RNTupleClusterBoundaries{currentBegin, currentEnd}); - } - return std::make_pair(adjustedClusters, nEntries); - } - - return std::make_pair(clustersInFile, nEntries); - }(); - const ROOT::NTupleSize_t nEntries = clustersAndEntries.second; - if (nEntries == ROOT::kInvalidNTupleIndex) - return; - const auto &clusters = clustersAndEntries.first; - ScheduleTaskForPageSource(pool, std::move(source), fileName, clusters, nEntries, lm, slotStack, processedEntries, - globalEntries); -} - -/// \brief Launch the task(s) for the input file index +/// \brief Run the computation graph with file-level tasks and cluster-level tail scheduling. /// -/// This function is used in MT scheduling to launch a task processing one input file. It is selected when -/// there is no global entry range. The task may be split into smaller tasks based on entry ranges aligned to cluster -/// boundaries, see ScheduleTaskForPageSource. -void ROOT::RDF::RNTupleDS::ScheduleTasksForFile(std::size_t fileIdx, ROOT::TThreadExecutor &pool, - ROOT::Detail::RDF::RLoopManager &lm, - ROOT::Internal::RSlotStack &slotStack, - std::atomic &processedEntries, - std::atomic &globalEntries) -{ - // Evaluate clusters (with local entry numbers) and number of entries for this file - const auto &fileName = fFileNames[fileIdx]; - auto source = CreatePageSource(fNTupleName, fFileNames[fileIdx]); - source->Attach(); - // TODO: improve MT scheduling such that not all page lists are required right away - source->LoadAllPageLists(); - const auto clustersAndEntries = [&]() { - const auto descGuard = source->GetSharedDescriptorGuard(); - return std::make_pair(ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()), descGuard->GetNEntries()); - }(); - const ROOT::NTupleSize_t nEntries = clustersAndEntries.second; - if (nEntries == 0) - return; - const auto &clusters = clustersAndEntries.first; - ScheduleTaskForPageSource(pool, std::move(source), fileName, clusters, nEntries, lm, slotStack, processedEntries, - globalEntries); -} - -/// \brief Run the computation graph by scheduling MT tasks -/// -/// The data source interacts with the loop manager and with a pool of threads (via TThreadExecutor) at the same time. -/// The function schedules one MT task per input file, eventually calling the parallel_for overload taking two input -/// arguments respectively respecting the Range and Body requirements. This means that each of the file-based tasks -/// can be split into various sub-tasks, each processing a set of consecutive clusters of the same file. Before -/// launching the file-based tasks, the function checks if the user has selected a global range of entries to restrict -/// processing to. If so, the function first computes the global entry offsets of all the input files, so they can be -/// matched against the user-provided global range. These will be used to adjust the list of clusters of the file before -/// spawning the corresponding task. +/// 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); - if (fGlobalEntryRange.has_value()) { - const auto globalFileOffsets = ComputeGlobalEntryOffsets(pool); - pool.Foreach( - [&](std::size_t fileIdx) { - ScheduleTasksForFileWithGlobalOffset(fileIdx, globalFileOffsets, pool, lm, slotStack, processedEntries, - globalEntries); - }, - ROOT::TSeq(fFileNames.size())); - } else { - pool.Foreach( - [&](std::size_t fileIdx) { - ScheduleTasksForFile(fileIdx, pool, lm, slotStack, processedEntries, globalEntries); - }, - ROOT::TSeq(fFileNames.size())); + std::vector globalFileOffsets; + if (fGlobalEntryRange) { + globalFileOffsets = ComputeGlobalEntryOffsets(pool); + assert(globalFileOffsets.size() == nFiles); } - if (fGlobalEntryRange.has_value()) { + 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: 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 7645f7026b85c..1542cff898fcb 100644 --- a/tree/dataframe/test/datasource_ntuple.cxx +++ b/tree/dataframe/test/datasource_ntuple.cxx @@ -338,6 +338,89 @@ TEST_F(RNTupleDSTest, ChainTailScheduling) 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 @@ -375,6 +458,14 @@ TEST(RNTupleDS, MoreFilesThanFileDescriptorLimit) 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