MakeRootTreeWriterSpec.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
#ifndef FRAMEWORK_MAKEROOTTREEWRITERSPEC_H
#define FRAMEWORK_MAKEROOTTREEWRITERSPEC_H
 
 
#include "DPLUtils/RootTreeWriter.h"
#include "Framework/InputSpec.h"
#include "Framework/ConfigParamRegistry.h"
#include "Framework/DataProcessorSpec.h"
#include "Framework/CallbackService.h"
#include "Framework/ControlService.h"
#include <algorithm>
#include <vector>
#include <string>
#include <stdexcept>
#include <variant>
#include <unordered_set>
#include <tuple>
 
namespace o2
{
namespace framework
{
 
class MakeRootTreeWriterSpec
{
 public:
  using WriterType = RootTreeWriter;
  enum struct TerminationPolicy {
    Process,
    Workflow,
  };
  const std::map<std::string, TerminationPolicy> TerminationPolicyMap = {
    {"process", TerminationPolicy::Process},
    {"workflow", TerminationPolicy::Workflow},
  };
 
  struct TerminationCondition {
    enum struct Action {
      DoProcessing,
      SkipProcessing,
    };
    using CheckProcessing = std::function<std::tuple<TerminationCondition::Action, bool>(framework::DataRef const&)>;
    using CheckReady = std::function<bool(o2::framework::DataRef const&)>;
 
    std::variant<std::monostate, CheckReady, CheckProcessing> check;
  };
 
  struct Preprocessor {
    using Process = std::function<void(ProcessingContext&)>;
 
    std::variant<std::monostate, Process> callback;
 
    constexpr operator bool()
    {
      return std::holds_alternative<Process>(callback);
    }
 
    void operator()(ProcessingContext& context)
    {
      if (std::holds_alternative<Process>(callback)) {
        std::get<Process>(callback)(context);
      }
    }
  };
 
  struct KeyExtractor {
    static std::string asString(InputSpec const& arg) { return arg.binding; }
  };
  template <typename T, typename _ = void>
  struct StringAssignable : public std::false_type {
  };
  template <typename T>
  struct StringAssignable<
    T,
    std::enable_if_t<std::is_same_v<std::decay_t<T>, char const*> ||
                     std::is_same_v<std::decay_t<T>, char*> ||
                     std::is_same_v<std::decay_t<T>, std::string>>> : public std::true_type {
  };
 
  template <typename T>
  struct BranchDefinition : public WriterType::BranchDef<T, InputSpec, KeyExtractor> {
    template <typename KeyType, typename Arg = const char*, std::enable_if_t<StringAssignable<Arg>::value, int> = 0>
    BranchDefinition(KeyType&& key, std::string _branchName, Arg _optionKey = "")
      : WriterType::BranchDef<T, InputSpec, KeyExtractor>(std::forward<KeyType>(key), _branchName), optionKey(_optionKey)
    {
    }
    template <typename KeyType, typename Arg = const char*, std::enable_if_t<StringAssignable<Arg>::value, int> = 0>
    BranchDefinition(KeyType&& key, std::string _branchName, int _nofBranches, Arg _optionKey = "")
      : WriterType::BranchDef<T, InputSpec, KeyExtractor>(std::forward<KeyType>(key), _branchName, _nofBranches), optionKey(_optionKey)
    {
    }
    template <typename KeyType, typename Arg, typename... Args, std::enable_if_t<StringAssignable<Arg>::value, int> = 0>
    BranchDefinition(KeyType key, std::string _branchName, Arg&& _optionKey, Args&&... args)
      : WriterType::BranchDef<T, InputSpec, KeyExtractor>(std::forward<KeyType>(key), _branchName, std::forward<Args>(args)...), optionKey(_optionKey)
    {
    }
    template <typename KeyType, typename Arg, typename... Args, std::enable_if_t<!StringAssignable<Arg>::value, int> = 0>
    BranchDefinition(KeyType key, std::string _branchName, Arg&& arg, Args&&... args)
      : WriterType::BranchDef<T, InputSpec, KeyExtractor>(std::forward<KeyType>(key), _branchName, std::forward<Arg>(arg), std::forward<Args>(args)...), optionKey()
    {
    }
    std::string optionKey = "";
  };
 
  // helper to define auxiliary inputs, i.e. inputs not connected to a branch
  struct AuxInputRoute {
    InputSpec mSpec;
    operator InputSpec() const
    {
      return mSpec;
    }
  };
 
  // helper to define tree attributes
  struct TreeAttributes {
    std::string name;
    std::string title = "";
  };
 
  // callback with signature void(TFile*, TTree*)
  using CustomClose = WriterType::CustomClose;
 
  MakeRootTreeWriterSpec() = delete;
 
  template <typename... Args>
  MakeRootTreeWriterSpec(const char* processName, Args&&... args) : mProcessName(processName)
  {
    parseConstructorArgs<0>(std::forward<Args>(args)...);
  }
 
  DataProcessorSpec operator()()
  {
    // this creates the workflow spec
    struct ProcessAttributes {
      // worker instance
      std::shared_ptr<WriterType> writer;
      // keys for branch name options
      std::vector<std::pair<std::string, std::string>> branchNameOptions;
      // number of events to be processed
      int nEvents = -1;
      // autosave every nEventsAutoSave events
      int nEventsAutoSave = -1;
      // starting with all inputs, every input which has been indicated 'ready' is removed
      std::unordered_set<std::string> activeInputs;
      // event counter
      int counter = 0;
      // indicate what to terminate upon ready: process or workflow
      TerminationPolicy terminationPolicy = TerminationPolicy::Process;
      // custom termination condition
      TerminationCondition terminationCondition;
      // custom preprocessor
      Preprocessor preprocessor;
      // the total number of served branches on the n inputs
      size_t nofBranches;
    };
    auto processAttributes = std::make_shared<ProcessAttributes>();
    processAttributes->writer = mWriter;
    processAttributes->branchNameOptions = mBranchNameOptions;
    processAttributes->terminationCondition = std::move(mTerminationCondition);
    processAttributes->preprocessor = std::move(mPreprocessor);
 
    // set the list of active inputs, every input which is indecated as 'ready' will be removed from the list
    // the process is ready if the list is empty
    for (auto const& input : mInputs) {
      processAttributes->activeInputs.emplace(input.binding);
    }
    processAttributes->nofBranches = mNofBranches;
 
    // the init function is returned to the DPL in order to init the process
    auto initFct = [processAttributes, TerminationPolicyMap = TerminationPolicyMap](InitContext& ic) {
      auto& branchNameOptions = processAttributes->branchNameOptions;
      auto filename = ic.options().get<std::string>("outfile");
      auto treename = ic.options().get<std::string>("treename");
      auto treetitle = ic.options().get<std::string>("treetitle");
      auto outdir = ic.options().get<std::string>("output-dir");
      processAttributes->nEvents = ic.options().get<int>("nevents");
      if (processAttributes->nEvents > 0 && processAttributes->activeInputs.size() != processAttributes->nofBranches) {
        LOG(warning) << "the n inputs serve in total m branches with n != m, this means that there will be data for\n"
                     << "different branches on the same input. Be aware that the --nevents option might lead to incomplete\n"
                     << "data in the output file as the number of processed input sets is counted";
      }
      processAttributes->nEventsAutoSave = ic.options().get<int>("autosave");
      try {
        processAttributes->terminationPolicy = TerminationPolicyMap.at(ic.options().get<std::string>("terminate"));
      } catch (std::out_of_range&) {
        throw std::invalid_argument(std::string("invalid termination policy: ") + ic.options().get<std::string>("terminate"));
      }
      if (filename.empty() || treename.empty()) {
        throw std::invalid_argument("output file name and tree name are mandatory options");
      }
      for (size_t branchIndex = 0; branchIndex < branchNameOptions.size(); branchIndex++) {
        // pair of key (first) - value (second)
        if (branchNameOptions[branchIndex].first.empty()) {
          continue;
        }
        auto branchName = ic.options().get<std::string>(branchNameOptions[branchIndex].first.c_str());
        processAttributes->writer->setBranchName(branchIndex, branchName.c_str());
      }
      if (!outdir.empty() && outdir != "none") {
        if (outdir.back() != '/') {
          outdir += '/';
        }
        filename = outdir + filename;
      }
      processAttributes->writer->init(filename.c_str(), treename.c_str(), treetitle.c_str());
      // the callback to be set as hook at stop of processing for the framework
      auto finishWriting = [processAttributes]() {
        processAttributes->writer->close();
      };
 
      ic.services().get<CallbackService>().set<CallbackService::Id::Stop>(finishWriting);
 
      auto processingFct = [processAttributes](ProcessingContext& pc) {
        auto& writer = processAttributes->writer;
        auto& terminationPolicy = processAttributes->terminationPolicy;
        auto& terminationCondition = processAttributes->terminationCondition;
        auto& preprocessor = processAttributes->preprocessor;
        auto& activeInputs = processAttributes->activeInputs;
        auto& counter = processAttributes->counter;
        auto& nEvents = processAttributes->nEvents;
        auto& nEventsAutoSave = processAttributes->nEventsAutoSave;
        if (writer->isClosed()) {
          return;
        }
 
        // if the termination condition contains function of type CheckProcessing, this is checked
        // before the processing, currently implemented logic:
        // - processing is skipped if this is indicated by at least one input
        // - if input is 'ready' it is removed from the list of active inputs
        auto checkProcessing = [&terminationCondition, &activeInputs](auto const& inputs) {
          bool doProcessing = true;
          if (std::holds_alternative<TerminationCondition::CheckProcessing>(terminationCondition.check)) {
            auto& check = std::get<TerminationCondition::CheckProcessing>(terminationCondition.check);
            for (auto const& ref : inputs) {
              auto iter = activeInputs.find(ref.spec->binding);
              auto [action, ready] = check(ref);
              if (action == TerminationCondition::Action::SkipProcessing) {
                // this condition tells us not to process the data any further
                doProcessing = false;
              }
              if (iter != activeInputs.end() && ready) {
                // this input is ready, remove from active inputs
                activeInputs.erase(iter);
              }
            }
          }
          return doProcessing;
        };
        // if the termination condition contains function of type CheckReady, this is checked
        // after the processing. All inputs marked 'ready' are removed from the list of active inputs,
        // process treated as 'ready' if no active inputs
        auto checkReady = [&terminationCondition, &activeInputs](auto const& inputs) {
          if (std::holds_alternative<TerminationCondition::CheckReady>(terminationCondition.check)) {
            auto& check = std::get<TerminationCondition::CheckReady>(terminationCondition.check);
            for (auto const& ref : inputs) {
              auto iter = activeInputs.find(ref.spec->binding);
              if (iter != activeInputs.end() && check(ref)) {
                // this input is ready, remove from active inputs
                activeInputs.erase(iter);
              }
            }
          }
          return activeInputs.size() == 0;
        };
 
        if (preprocessor) {
          preprocessor(pc);
        }
        if (checkProcessing(pc.inputs())) {
          (*writer)(pc.inputs());
          counter = counter + 1;
        }
 
        if ((nEvents >= 0 && counter == nEvents) || checkReady(pc.inputs())) {
          writer->close();
          pc.services().get<ControlService>().readyToQuit(terminationPolicy == TerminationPolicy::Workflow ? QuitRequest::All : QuitRequest::Me);
        } else if (nEventsAutoSave > 0 && counter && (counter % nEventsAutoSave) == 0) {
          writer->autoSave();
        }
      };
 
      return processingFct;
    };
 
    Options options{
      // default options
      {"outfile", VariantType::String, mDefaultFileName.c_str(), {"Name of the output file"}},
      {"output-dir", VariantType::String, mDefaultDir.c_str(), {"Output directory"}},
      {"treename", VariantType::String, mDefaultTreeName.c_str(), {"Name of tree"}},
      {"treetitle", VariantType::String, mDefaultTreeTitle.c_str(), {"Title of tree"}},
      {"nevents", VariantType::Int, mDefaultNofEvents, {"Number of events to execute"}},
      {"autosave", VariantType::Int, mDefaultAutoSave, {"Autosave after number of events"}},
      {"terminate", VariantType::String, mDefaultTerminationPolicy.c_str(), {"Terminate the 'process' or 'workflow'"}},
    };
    for (size_t branchIndex = 0; branchIndex < mBranchNameOptions.size(); branchIndex++) {
      // adding option definitions for those ones defined in the branch definition
      if (mBranchNameOptions[branchIndex].first.empty()) {
        continue;
      }
      options.push_back(ConfigParamSpec(mBranchNameOptions[branchIndex].first.c_str(),  // option key
                                        VariantType::String,                            // option argument type
                                        mBranchNameOptions[branchIndex].second.c_str(), // default branch name
                                        {"configurable branch name"}                    // help message
                                        ));
    }
 
    mInputRoutes.insert(mInputRoutes.end(), mInputs.begin(), mInputs.end());
    return DataProcessorSpec{
      // processing spec generated from the class configuartion
      mProcessName.c_str(),   // name of the process
      mInputRoutes,           // list of inputs
      Outputs{},              // no outputs
      AlgorithmSpec(initFct), // return the init function
      std::move(options),     // processor options
    };
  }
 
 private:
  template <size_t N, typename... Args>
  void parseConstructorArgs(const char* name, Args&&... args)
  {
    static_assert(N == 0, "wrong argument order, default file and tree options must come before branch specs");
    // this can be called twice, the first time we set the default file name
    // and if we are here for a second time, we set the default tree name
    if (mDefaultFileName.empty()) {
      mDefaultFileName = name;
    } else {
      mDefaultTreeName = name;
      mDefaultTreeTitle = name;
    }
 
    parseConstructorArgs<N>(std::forward<Args>(args)...);
  }
 
  template <size_t N, typename... Args>
  void parseConstructorArgs(int arg, Args&&... args)
  {
    static_assert(N == 0, "wrong argument order, default file and tree options must come before branch specs");
    if (mNIntArgCounter == 0) {
      mDefaultNofEvents = arg;
    } else if (mNIntArgCounter == 1) {
      mDefaultAutoSave = arg;
    } else {
      throw std::logic_error("Too many integer arguments in the constructor");
    }
    mNIntArgCounter++;
    parseConstructorArgs<N>(std::forward<Args>(args)...);
  }
 
  template <size_t N, typename... Args>
  void parseConstructorArgs(TerminationPolicy arg, Args&&... args)
  {
    static_assert(N == 0, "wrong argument order, default file and tree, and all options must come before branch specs");
    for (const auto& policy : TerminationPolicyMap) {
      if (policy.second == arg) {
        mDefaultTerminationPolicy = policy.first;
        return parseConstructorArgs<N>(std::forward<Args>(args)...);
      }
    }
    // here we only get if the enum and map definitions are inconsistent
    throw std::logic_error("Internal mismatch of policy ids and keys");
  }
 
  template <size_t N, typename... Args>
  void parseConstructorArgs(TerminationCondition&& arg, Args&&... args)
  {
    static_assert(N == 0, "wrong argument order, default file and tree, and all options must come before branch specs");
    mTerminationCondition = std::move(arg);
    parseConstructorArgs<N>(std::forward<Args>(args)...);
  }
 
  template <size_t N, typename... Args>
  void parseConstructorArgs(Preprocessor&& arg, Args&&... args)
  {
    static_assert(N == 0, "wrong argument order, default file and tree, and all options must come before branch specs");
    mPreprocessor = std::move(arg);
    parseConstructorArgs<N>(std::forward<Args>(args)...);
  }
 
  template <size_t N, typename... Args>
  void parseConstructorArgs(AuxInputRoute&& aux, Args&&... args)
  {
    mInputRoutes.emplace_back(aux);
    parseConstructorArgs<N>(std::forward<Args>(args)...);
  }
 
  template <size_t N, typename... Args>
  void parseConstructorArgs(CustomClose&& callback, Args&&... args)
  {
    mCustomClose = callback;
    parseConstructorArgs<N>(std::forward<Args>(args)...);
  }
 
  template <size_t N, typename... Args>
  void parseConstructorArgs(TreeAttributes&& att, Args&&... args)
  {
    mDefaultTreeName = att.name;
    mDefaultTreeTitle = att.title.empty() ? att.name : att.title;
    parseConstructorArgs<N>(std::forward<Args>(args)...);
  }
 
  template <size_t N, typename T, typename... Args>
  void parseConstructorArgs(BranchDefinition<T>&& def, Args&&... args)
  {
    if (def.nofBranches > 0) {
      // number of branches set to 0 will skip the definition, this allows to
      // dynamically disable branches, while all possible definitions can
      // be specified at compile time
      mInputs.insert(mInputs.end(), def.keys.begin(), def.keys.end());
      mBranchNameOptions.emplace_back(def.optionKey, def.branchName);
      mNofBranches += def.nofBranches;
    } else {
      // insert an empty placeholder
      mBranchNameOptions.emplace_back("", "");
    }
    parseConstructorArgs<N + 1>(std::forward<Args>(args)...);
    if constexpr (N == 0) {
      mWriter = std::make_shared<WriterType>(nullptr, nullptr, mCustomClose, std::forward<BranchDefinition<T>>(def), std::forward<Args>(args)...);
    }
  }
 
  // this terminates the argument parsing
  template <size_t N>
  void parseConstructorArgs()
  {
  }
 
  std::shared_ptr<WriterType> mWriter;
  std::string mProcessName;
  std::vector<InputSpec> mInputs;
  std::vector<InputSpec> mInputRoutes;
  std::vector<std::pair<std::string, std::string>> mBranchNameOptions;
  std::string mDefaultFileName;
  std::string mDefaultTreeName;
  std::string mDefaultTreeTitle;
  std::string mDefaultDir = "none";
  int mDefaultNofEvents = -1;
  int mDefaultAutoSave = -1;
  std::string mDefaultTerminationPolicy = "process";
  TerminationCondition mTerminationCondition;
  Preprocessor mPreprocessor;
  size_t mNofBranches = 0;
  int mNIntArgCounter = 0;
  CustomClose mCustomClose;
};
} // namespace framework
} // namespace o2
 
#endif // FRAMEWORK_MAKEROOTTREEWRITERSPEC_H