/* @(#)root/multiproc:$Id$ */
// Author: Enrico Guiraud July 2015
// Modified: G Ganis Jan 2017
/*************************************************************************
* Copyright (C) 1995-2000, Rene Brun and Fons Rademakers. *
* All rights reserved. *
* *
* For the licensing terms see $ROOTSYS/LICENSE. *
* For the list of contributors see $ROOTSYS/README/CREDITS. *
*************************************************************************/
#ifndef ROOT_TTreeProcessorMP
#define ROOT_TTreeProcessorMP
#include "MPCode.h"
#include "MPSendRecv.h"
#include "PoolUtils.h"
#include "TChain.h"
#include "TChainElement.h"
#include "TError.h"
#include "TFileCollection.h"
#include "TFileInfo.h"
#include "THashList.h"
#include "TMPClient.h"
#include "TMPWorkerTree.h"
#include "TSelector.h"
#include "TTreeReader.h"
#include <algorithm> //std::generate
#include <numeric> //std::iota
#include <string>
#include <type_traits> //std::result_of, std::enable_if
#include <functional> //std::reference_wrapper
#include <vector>
namespace ROOT {
class TTreeProcessorMP : private TMPClient {
public:
explicit TTreeProcessorMP(unsigned nWorkers = 0); //default number of workers is the number of processors
~TTreeProcessorMP() = default;
//it doesn't make sense for a TTreeProcessorMP to be copied
TTreeProcessorMP(const TTreeProcessorMP &) = delete;
TTreeProcessorMP &operator=(const TTreeProcessorMP &) = delete;
/// \brief Process a TTree dataset with a functor
/// \tparam F functor returning a pointer to TObject or inheriting classes and
/// taking a TTreeReader& (both enforced at compile-time)
///
/// Dataset definition:
/// \param[in] fileNames vector of strings with the paths of the files with the TTree to process
/// \param[in] fileName string with the path of the files with the TTree to process
/// \param[in] collection TFileCollection with the files with the TTree to process
/// \param[in] chain TChain with the files with the TTree to process
/// \param[in] tree TTree to process
///
/// \param[in] entries TEntryList to filter the dataset
/// \param[in] treeName Name of the TTree to process
/// \param[in] nToProcess Number of entries to process (0 means all)
/// \param[in] jFirst First entry to process (0 means the first of the first file)
///
template<class F> auto Process(const std::vector<std::string>& fileNames, F procFunc, TEntryList &entries,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
template<class F> auto Process(const std::string& fileName, F procFunc, TEntryList &entries,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
template<class F> auto Process(TFileCollection& collection, F procFunc, TEntryList &entries,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
template<class F> auto Process(TChain& chain, F procFunc, TEntryList &entries,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
template<class F> auto Process(TTree& tree, F procFunc, TEntryList &entries,
ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
/// \brief Process a TTree dataset with a functor: version without entry list
/// \tparam F functor returning a pointer to TObject or inheriting classes and
/// taking a TTreeReader& (both enforced at compile-time)
///
/// Dataset definition:
/// \param[in] fileNames vector of strings with the paths of the files with the TTree to process
/// \param[in] fileName string with the path of the files with the TTree to process
/// \param[in] collection TFileCollection with the files with the TTree to process
/// \param[in] chain TChain with the files with the TTree to process
/// \param[in] tree TTree to process
///
/// \param[in] treeName Name of the TTree to process
/// \param[in] nToProcess Number of entries to process (0 means all)
/// \param[in] jFirst First entry to process (0 means the first of the first file)
///
template<class F> auto Process(const std::vector<std::string>& fileNames, F procFunc,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
template<class F> auto Process(const std::string& fileName, F procFunc,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
template<class F> auto Process(TFileCollection& files, F procFunc,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
template<class F> auto Process(TChain& files, F procFunc,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
template<class F> auto Process(TTree& tree, F procFunc, ULong64_t nToProcess = 0, ULong64_t jFirst = 0)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
/// \brief Process a TTree dataset with a selector
///
/// Dataset definition:
/// \param[in] fileNames vector of strings with the paths of the files with the TTree to process
/// \param[in] fileName string with the path of the files with the TTree to process
/// \param[in] collection TFileCollection with the files with the TTree to process
/// \param[in] chain TChain with the files with the TTree to process
/// \param[in] tree TTree to process
///
/// \param[in] selector Instance of TSelector to be applied to the dataset
/// \param[in] entries TEntryList to filter the dataset
/// \param[in] treeName Name of the TTree to process
/// \param[in] nToProcess Number of entries to process (0 means all)
/// \param[in] jFirst First entry to process (0 means the first of the first file)
///
// these versions require a TSelector
TList* Process(const std::vector<std::string>& fileNames, TSelector& selector, TEntryList &entries,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
TList* Process(const std::string &fileName, TSelector& selector, TEntryList &entries,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
TList* Process(TFileCollection& files, TSelector& selector, TEntryList &entries,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
TList* Process(TChain& files, TSelector& selector, TEntryList &entries,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
TList* Process(TTree& tree, TSelector& selector, TEntryList &entries,
ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
/// \brief Process a TTree dataset with a selector: version without entry list
///
/// Dataset definition:
/// \param[in] fileNames vector of strings with the paths of the files with the TTree to process
/// \param[in] fileName string with the path of the files with the TTree to process
/// \param[in] collection TFileCollection with the files with the TTree to process
/// \param[in] chain TChain with the files with the TTree to process
/// \param[in] tree TTree to process
///
/// \param[in] selector Instance of TSelector to be applied to the dataset
/// \param[in] treeName Name of the TTree to process
/// \param[in] nToProcess Number of entries to process (0 means all)
/// \param[in] jFirst First entry to process (0 means the first of the first file)
///
// these versions require a TSelector
TList* Process(const std::vector<std::string>& fileNames, TSelector& selector,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
TList* Process(const std::string &fileName, TSelector& selector,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
TList* Process(TFileCollection& files, TSelector& selector,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
TList* Process(TChain& files, TSelector& selector,
const std::string& treeName = "", ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
TList* Process(TTree& tree, TSelector& selector, ULong64_t nToProcess = 0, ULong64_t jFirst = 0);
void SetNWorkers(unsigned n) { TMPClient::SetNWorkers(n); }
unsigned GetNWorkers() const { return TMPClient::GetNWorkers(); }
private:
template<class T> void Collect(std::vector<T> &reslist);
template<class T> void HandlePoolCode(MPCodeBufPair &msg, TSocket *sender, std::vector<T> &reslist);
void FixLists(std::vector<TObject*> &lists);
void Reset();
void ReplyToIdle(TSocket *s);
unsigned fNProcessed; ///< number of arguments already passed to the workers
unsigned fNToProcess; ///< total number of arguments to pass to the workers
/// A collection of the types of tasks that TTreeProcessorMP can execute.
/// It is used to interpret in the right way and properly reply to the
/// messages received (see, for example, TTreeProcessorMP::HandleInput)
enum class ETask : unsigned char {
kNoTask, ///< no task is being executed
kProcByRange, ///< a Process method is being executed and each worker will process a certain range of each file
kProcByFile ///< a Process method is being executed and each worker will process a different file
};
ETask fTaskType = ETask::kNoTask; ///< the kind of task that is being executed, if any
};
template<class F>
auto TTreeProcessorMP::Process(const std::vector<std::string>& fileNames, F procFunc, TEntryList &entries,
const std::string& treeName, ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
using retType = typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
static_assert(std::is_constructible<TObject*, retType>::value,
"procFunc must return a pointer to a class inheriting from TObject,"
" and must take a reference to TTreeReader as the only argument");
// Warn for yet unimplemented functionality
if (jFirst > 0) {
Warning("Process", "support for generic 'first entry' (jFirst > 0) not implemented yet - ignoring");
jFirst = 0;
}
//prepare environment
Reset();
unsigned nWorkers = GetNWorkers();
// Check th entry list
TEntryList *elist = (entries.IsValid()) ? &entries : nullptr;
//fork
TMPWorkerTreeFunc<F> worker(procFunc, fileNames, elist, treeName, nWorkers, nToProcess, jFirst);
bool ok = Fork(worker);
if(!ok) {
Error("TTreeProcessorMP::Process", "[E][C] Could not fork. Aborting operation.");
return nullptr;
}
if(fileNames.size() < nWorkers) {
//TTree entry granularity. For each file, we divide entries equally between workers
fTaskType = ETask::kProcByRange;
//Tell workers to start processing entries
fNToProcess = nWorkers*fileNames.size(); //this is the total number of ranges that will be processed by all workers cumulatively
std::vector<unsigned> args(nWorkers);
std::iota(args.begin(), args.end(), 0);
fNProcessed = Broadcast(MPCode::kProcRange, args);
if(fNProcessed < nWorkers)
Error("TTreeProcessorMP::Process", "[E][C] There was an error while sending tasks to workers. Some entries might not be processed.");
} else {
//file granularity. each worker processes one whole file as a single task
fTaskType = ETask::kProcByFile;
fNToProcess = fileNames.size();
std::vector<unsigned> args(nWorkers);
std::iota(args.begin(), args.end(), 0);
fNProcessed = Broadcast(MPCode::kProcFile, args);
if(fNProcessed < nWorkers)
Error("TTreeProcessorMP::Process", "[E][C] There was an error while sending tasks to workers. Some entries might not be processed.");
}
//collect results, distribute new tasks
std::vector<TObject*> reslist;
Collect(reslist);
//merge
PoolUtils::ReduceObjects<TObject *> redfunc;
auto res = redfunc(reslist);
//clean-up and return
ReapWorkers();
fTaskType = ETask::kNoTask;
return static_cast<retType>(res);
}
template<class F>
auto TTreeProcessorMP::Process(const std::string& fileName, F procFunc, TEntryList &entries,
const std::string& treeName, ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
std::vector<std::string> singleFileName(1, fileName);
return Process(singleFileName, procFunc, entries, treeName, nToProcess, jFirst);
}
template<class F>
auto TTreeProcessorMP::Process(TFileCollection& files, F procFunc, TEntryList &entries,
const std::string& treeName, ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
std::vector<std::string> fileNames(files.GetNFiles());
unsigned count = 0;
for(auto f : *static_cast<THashList*>(files.GetList()))
fileNames[count++] = static_cast<TFileInfo*>(f)->GetCurrentUrl()->GetUrl();
return Process(fileNames, procFunc, entries, treeName, nToProcess, jFirst);
}
template<class F>
auto TTreeProcessorMP::Process(TChain& files, F procFunc, TEntryList &entries,
const std::string& treeName, ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
TObjArray* filelist = files.GetListOfFiles();
std::vector<std::string> fileNames(filelist->GetEntries());
unsigned count = 0;
for(auto f : *filelist)
fileNames[count++] = f->GetTitle();
return Process(fileNames, procFunc, entries, treeName, nToProcess, jFirst);
}
template<class F>
auto TTreeProcessorMP::Process(TTree& tree, F procFunc, TEntryList &entries,
ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
using retType = typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type;
static_assert(std::is_constructible<TObject*, retType>::value, "procFunc must return a pointer to a class inheriting from TObject, and must take a reference to TTreeReader as the only argument");
// Warn for yet unimplemented functionality
if (jFirst > 0) {
Warning("Process", "support for generic 'first entry' (jFirst > 0) not implemented yet - ignoring");
jFirst = 0;
}
//prepare environment
Reset();
unsigned nWorkers = GetNWorkers();
// Check th entry list
TEntryList *elist = (entries.IsValid()) ? &entries : nullptr;
//fork
TMPWorkerTreeFunc<F> worker(procFunc, &tree, elist, nWorkers, nToProcess, jFirst);
bool ok = Fork(worker);
if(!ok) {
Error("TTreeProcessorMP::Process", "[E][C] Could not fork. Aborting operation.");
return nullptr;
}
//divide entries equally between workers
fTaskType = ETask::kProcByRange;
//tell workers to start processing entries
fNToProcess = nWorkers; //this is the total number of ranges that will be processed by all workers cumulatively
std::vector<unsigned> args(nWorkers);
std::iota(args.begin(), args.end(), 0);
fNProcessed = Broadcast(MPCode::kProcTree, args);
if(fNProcessed < nWorkers)
Error("TTreeProcessorMP::Process", "[E][C] There was an error while sending tasks to workers. Some entries might not be processed.");
//collect results, distribute new tasks
std::vector<TObject*> reslist;
Collect(reslist);
//merge
PoolUtils::ReduceObjects<TObject *> redfunc;
auto res = redfunc(reslist);
//clean-up and return
ReapWorkers();
fTaskType = ETask::kNoTask;
return static_cast<retType>(res);
}
///
/// No TEntryList versions of generic processor
///
template<class F>
auto TTreeProcessorMP::Process(const std::vector<std::string>& fileNames, F procFunc,
const std::string& treeName, ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
TEntryList noelist;
return Process(fileNames, procFunc, noelist, treeName, nToProcess, jFirst);
}
template<class F>
auto TTreeProcessorMP::Process(const std::string& fileName, F procFunc,
const std::string& treeName, ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
TEntryList noelist;
return Process(fileName, procFunc, noelist, treeName, nToProcess, jFirst);
}
template<class F>
auto TTreeProcessorMP::Process(TFileCollection& files, F procFunc,
const std::string& treeName, ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
TEntryList noelist;
return Process(files, procFunc, noelist, treeName, nToProcess, jFirst);
}
template<class F>
auto TTreeProcessorMP::Process(TChain& files, F procFunc,
const std::string& treeName, ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
TEntryList noelist;
return Process(files, procFunc, noelist, treeName, nToProcess, jFirst);
}
template<class F>
auto TTreeProcessorMP::Process(TTree& tree, F procFunc,
ULong64_t nToProcess, ULong64_t jFirst)
-> typename std::result_of<F(std::reference_wrapper<TTreeReader>)>::type
{
TEntryList noelist;
return Process(tree, procFunc, noelist, nToProcess, jFirst);
}
//////////////////////////////////////////////////////////////////////////
/// Handle message and reply to the worker
template<class T>
void TTreeProcessorMP::HandlePoolCode(MPCodeBufPair &msg, TSocket *s, std::vector<T> &reslist)
{
unsigned code = msg.first;
if (code == MPCode::kIdling) {
ReplyToIdle(s);
} else if(code == MPCode::kProcResult) {
if(msg.second != nullptr)
reslist.push_back(std::move(ReadBuffer<T>(msg.second.get())));
MPSend(s, MPCode::kShutdownOrder);
} else if(code == MPCode::kProcError) {
const char *str = ReadBuffer<const char*>(msg.second.get());
Error("TTreeProcessorMP::HandlePoolCode", "[E][C] a worker encountered an error: %s\n"
"Continuing execution ignoring these entries.", str);
ReplyToIdle(s);
delete [] str;
} else {
// UNKNOWN CODE
Error("TTreeProcessorMP::HandlePoolCode", "[W][C] unknown code received from server. code=%d", code);
}
}
//////////////////////////////////////////////////////////////////////////
/// Listen for messages sent by the workers and call the appropriate handler function.
/// TTreeProcessorMP::HandlePoolCode is called on messages with a code < 1000 and
/// TMPClient::HandleMPCode is called on messages with a code >= 1000.
template<class T>
void TTreeProcessorMP::Collect(std::vector<T> &reslist)
{
TMonitor &mon = GetMonitor();
mon.ActivateAll();
while (mon.GetActive() > 0) {
TSocket *s = mon.Select();
MPCodeBufPair msg = MPRecv(s);
if (msg.first == MPCode::kRecvError) {
Error("TTreeProcessorMP::Collect", "[E][C] Lost connection to a worker");
Remove(s);
} else if (msg.first < 1000)
HandlePoolCode(msg, s, reslist);
else
HandleMPCode(msg, s);
}
}
} // ROOT namespace
#endif