Utilities/plugins/OrbitNanoAODOutputModule.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:     L1TriggerScouting/Utilities
0004 // Class  :     OrbitNanoAODOutputModule
0005 //
0006 // Implementation:
0007 //     Adapt from NanoAODOutputModule for OrbitFlatTable
0008 //     This handles rotating OrbitCollection to Event
0009 //
0010 //
0011 // Author original version: Giovanni Petrucciani
0012 //        adapted by Patin Inkaew
0013
0014 // system include files
0015 #include <algorithm>
0016 #include <memory>
0017
0018 #include "Compression.h"
0019 #include "TFile.h"
0020 #include "TObjString.h"
0021 #include "TROOT.h"
0022 #include "TTree.h"
0023 #include <string>
0024
0025 // user include files
0026 #include "FWCore/Framework/interface/one/OutputModule.h"
0027 #include "FWCore/Framework/interface/RunForOutput.h"
0028 #include "FWCore/Framework/interface/LuminosityBlockForOutput.h"
0029 #include "FWCore/Framework/interface/EventForOutput.h"
0030 #include "FWCore/ServiceRegistry/interface/Service.h"
0031 #include "FWCore/Framework/interface/MakerMacros.h"
0032 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0033 #include "FWCore/MessageLogger/interface/JobReport.h"
0034 #include "FWCore/Utilities/interface/GlobalIdentifier.h"
0035 #include "FWCore/Utilities/interface/Digest.h"
0036 #include "IOPool/Provenance/interface/CommonProvenanceFiller.h"
0037 #include "DataFormats/Provenance/interface/BranchType.h"
0038 #include "DataFormats/Provenance/interface/BranchDescription.h"
0039 #include "DataFormats/Provenance/interface/ProcessHistoryRegistry.h"
0040 #include "DataFormats/NanoAOD/interface/FlatTable.h"
0041 #include "L1TriggerScouting/Utilities/plugins/OrbitTableOutputBranches.h"
0042 #include "L1TriggerScouting/Utilities/plugins/SelectedBxTableOutputBranches.h"
0043
0044 #include "oneapi/tbb/task_arena.h"
0045
0046 class OrbitNanoAODOutputModule : public edm::one::OutputModule<> {
0047 public:
0048   OrbitNanoAODOutputModule(edm::ParameterSet const& pset);
0049   ~OrbitNanoAODOutputModule() override;
0050
0051   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0052
0053 private:
0054   void write(edm::EventForOutput const& e) override;
0055   void writeLuminosityBlock(edm::LuminosityBlockForOutput const&) override;
0056   void writeRun(edm::RunForOutput const&) override;
0057   bool isFileOpen() const override;
0058   void openFile(edm::FileBlock const&) override;
0059   void reallyCloseFile() override;
0060
0061   std::string m_fileName;
0062   std::string m_logicalFileName;
0063   int m_compressionLevel;
0064   int m_eventsSinceFlush{0};
0065   std::string m_compressionAlgorithm;
0066   bool m_skipEmptyBXs;
0067   bool m_writeProvenance;
0068   bool m_fakeName;  //crab workaround, remove after crab is fixed
0069   int m_autoFlush;
0070   edm::ProcessHistoryRegistry m_processHistoryRegistry;
0071   edm::JobReport::Token m_jrToken;
0072   std::unique_ptr<TFile> m_file;
0073   std::unique_ptr<TTree> m_tree, m_lumiTree, m_runTree, m_metaDataTree, m_parameterSetsTree;
0074
0075   static constexpr int m_firstFlush{1000};
0076
0077   class CommonEventBranches {
0078   public:
0079     void branch(TTree& tree) {
0080       tree.Branch("run", &m_run, "run/i");
0081       tree.Branch("luminosityBlock", &m_luminosityBlock, "luminosityBlock/i");
0082       tree.Branch("bunchCrossing", &m_bunchCrossing, "bunchCrossing/i");
0083       tree.Branch("orbitNumber", &m_orbitNumber, "orbitNumber/i");
0084     }
0085     void fill(const edm::EventAuxiliary& aux) {
0086       m_run = aux.id().run();
0087       m_luminosityBlock = aux.id().luminosityBlock();
0088       m_orbitNumber = aux.id().event();  // in L1Scouting, one processing event is one orbit
0089     }
0090     void setBx(unsigned bx) { m_bunchCrossing = bx; }
0091
0092   private:
0093     UInt_t m_run;
0094     UInt_t m_luminosityBlock;
0095     UInt_t m_bunchCrossing;
0096     UInt_t m_orbitNumber;
0097   } m_commonBranches;
0098
0099   class CommonLumiBranches {
0100   public:
0101     void branch(TTree& tree) {
0102       tree.Branch("run", &m_run, "run/i");
0103       tree.Branch("luminosityBlock", &m_luminosityBlock, "luminosityBlock/i");
0104       tree.Branch("nOrbits", &m_orbits, "nOrbits/i");
0105     }
0106     void fill(const edm::LuminosityBlockID& id, unsigned int nOrbits) {
0107       m_run = id.run();
0108       m_luminosityBlock = id.value();
0109       m_orbits = nOrbits;
0110     }
0111
0112   private:
0113     UInt_t m_run;
0114     UInt_t m_luminosityBlock;
0115     UInt_t m_orbits;
0116   } m_commonLumiBranches;
0117
0118   class CommonRunBranches {
0119   public:
0120     void branch(TTree& tree) { tree.Branch("run", &m_run, "run/i"); }
0121     void fill(const edm::RunID& id) { m_run = id.run(); }
0122
0123   private:
0124     UInt_t m_run;
0125   } m_commonRunBranches;
0126
0127   std::vector<OrbitTableOutputBranches> m_tables;
0128   std::vector<SelectedBxTableOutputBranches> m_selbxs;
0129   unsigned int m_nOrbits;
0130
0131   std::vector<std::pair<std::string, edm::EDGetToken>> m_nanoMetadata;
0132
0133   edm::EDGetTokenT<std::vector<unsigned>> m_bxMaskToken;
0134   std::vector<unsigned> allBXs_;
0135 };
0136
0137 //
0138 // constants, enums and typedefs
0139 //
0140
0141 //
0142 // static data member definitions
0143 //
0144
0145 //
0146 // constructors and destructor
0147 //
0148 OrbitNanoAODOutputModule::OrbitNanoAODOutputModule(edm::ParameterSet const& pset)
0149     : edm::one::OutputModuleBase::OutputModuleBase(pset),
0150       edm::one::OutputModule<>(pset),
0151       m_fileName(pset.getUntrackedParameter<std::string>("fileName")),
0152       m_logicalFileName(pset.getUntrackedParameter<std::string>("logicalFileName")),
0153       m_compressionLevel(pset.getUntrackedParameter<int>("compressionLevel")),
0154       m_compressionAlgorithm(pset.getUntrackedParameter<std::string>("compressionAlgorithm")),
0155       m_skipEmptyBXs(pset.getParameter<bool>("skipEmptyBXs")),
0156       m_writeProvenance(pset.getUntrackedParameter<bool>("saveProvenance", true)),
0157       m_fakeName(pset.getUntrackedParameter<bool>("fakeNameForCrab", false)),
0158       m_autoFlush(pset.getUntrackedParameter<int>("autoFlush", -10000000)),
0159       m_processHistoryRegistry(),
0160       m_nOrbits(0) {
0161   edm::InputTag bxMask = pset.getParameter<edm::InputTag>("selectedBx");
0162   if (!bxMask.label().empty()) {
0163     m_bxMaskToken = consumes<std::vector<unsigned>>(bxMask);
0164   } else {
0165     allBXs_.resize(l1ScoutingRun3::OrbitFlatTable::NBX);
0166     for (unsigned int i = 0; i < l1ScoutingRun3::OrbitFlatTable::NBX; ++i) {
0167       allBXs_[i] = i + 1;
0168     }
0169   }
0170 }
0171
0172 OrbitNanoAODOutputModule::~OrbitNanoAODOutputModule() {}
0173
0174 void OrbitNanoAODOutputModule::write(edm::EventForOutput const& iEvent) {
0175   //Get data from 'e' and write it to the file
0176   edm::Service<edm::JobReport> jr;
0177   jr->eventWrittenToFile(m_jrToken, iEvent.id().run(), iEvent.id().event());
0178   m_nOrbits++;
0179
0180   if (m_autoFlush) {
0181     int64_t events = m_tree->GetEntriesFast();
0182     if (events == m_firstFlush) {
0183       m_tree->FlushBaskets();
0184       float maxMemory;
0185       if (m_autoFlush > 0) {
0186         // Estimate the memory we'll be using at the first full flush by
0187         // linearly scaling the number of events.
0188         float percentClusterDone = m_firstFlush / static_cast<float>(m_autoFlush);
0189         maxMemory = static_cast<float>(m_tree->GetTotBytes()) / percentClusterDone;
0190       } else if (m_tree->GetZipBytes() == 0) {
0191         maxMemory = 100 * 1024 * 1024;  // Degenerate case of no information in the tree; arbitrary value
0192       } else {
0193         // Estimate the memory we'll be using by scaling the current compression ratio.
0194         float cxnRatio = m_tree->GetTotBytes() / static_cast<float>(m_tree->GetZipBytes());
0195         maxMemory = -m_autoFlush * cxnRatio;
0196         float percentBytesDone = -m_tree->GetZipBytes() / static_cast<float>(m_autoFlush);
0197         m_autoFlush = m_firstFlush / percentBytesDone;
0198       }
0199       //std::cout << "OptimizeBaskets: total bytes " << m_tree->GetTotBytes() << std::endl;
0200       //std::cout << "OptimizeBaskets: zip bytes " << m_tree->GetZipBytes() << std::endl;
0201       //std::cout << "OptimizeBaskets: autoFlush " << m_autoFlush << std::endl;
0202       //std::cout << "OptimizeBaskets: maxMemory " << static_cast<uint32_t>(maxMemory) << std::endl;
0203       //m_tree->OptimizeBaskets(static_cast<uint32_t>(maxMemory), 1, "d");
0204       m_tree->OptimizeBaskets(static_cast<uint32_t>(maxMemory), 1, "");
0205     }
0206     if (m_eventsSinceFlush == m_autoFlush) {
0207       m_tree->FlushBaskets();
0208       m_eventsSinceFlush = 0;
0209     }
0210     m_eventsSinceFlush++;
0211   }
0212
0213   m_commonBranches.fill(iEvent.eventAuxiliary());
0214   // fill all tables, starting from main tables and then doing extension tables
0215   for (unsigned int extensions = 0; extensions <= 1; ++extensions) {
0216     for (auto& t : m_tables) {
0217       t.beginFill(iEvent, *m_tree, extensions);
0218     }
0219   }
0220   // get m_table, read parameters, book branches, etc.
0221   for (auto& t : m_selbxs) {
0222     t.beginFill(iEvent, *m_tree);
0223   }
0224   // get a lust of selected BXs to be filled
0225   const std::vector<unsigned>* selbx = &allBXs_;
0226   if (!m_bxMaskToken.isUninitialized()) {
0227     edm::Handle<std::vector<unsigned>> handle;
0228     iEvent.getByToken(m_bxMaskToken, handle);
0229     selbx = &*handle;
0230   }
0231   // convert from orbit as event to collision as event
0232   tbb::this_task_arena::isolate([&] {
0233     for (unsigned bx : *selbx) {
0234       if (m_skipEmptyBXs) {
0235         bool empty = true;
0236         for (auto& t : m_tables) {
0237           if (t.hasBx(bx)) {
0238             empty = false;
0239             break;
0240           }
0241         }
0242         if (empty) {
0243           continue;
0244         }
0245       }
0246
0247       m_commonBranches.setBx(bx);
0248       for (auto& t : m_tables) {
0249         t.fillBx(bx);
0250       }
0251       for (auto& t : m_selbxs) {
0252         t.fillBx(bx);
0253       }
0254       m_tree->Fill();
0255     }  // bx loop
0256   });
0257
0258   // set m_table to nullptr
0259   for (auto& t : m_tables) {
0260     t.endFill();
0261   }
0262   for (auto& t : m_selbxs) {
0263     t.endFill();
0264   }
0265   m_processHistoryRegistry.registerProcessHistory(iEvent.processHistory());
0266 }
0267
0268 void OrbitNanoAODOutputModule::writeLuminosityBlock(edm::LuminosityBlockForOutput const& iLumi) {
0269   edm::Service<edm::JobReport> jr;
0270   jr->reportLumiSection(m_jrToken, iLumi.id().run(), iLumi.id().value());
0271
0272   m_commonLumiBranches.fill(iLumi.id(), m_nOrbits);
0273
0274   tbb::this_task_arena::isolate([&] { m_lumiTree->Fill(); });
0275
0276   m_processHistoryRegistry.registerProcessHistory(iLumi.processHistory());
0277
0278   m_nOrbits = 0;
0279 }
0280
0281 void OrbitNanoAODOutputModule::writeRun(edm::RunForOutput const& iRun) {
0282   edm::Service<edm::JobReport> jr;
0283   jr->reportRunNumber(m_jrToken, iRun.id().run());
0284
0285   m_commonRunBranches.fill(iRun.id());
0286
0287   tbb::this_task_arena::isolate([&] { m_runTree->Fill(); });
0288
0289   m_processHistoryRegistry.registerProcessHistory(iRun.processHistory());
0290 }
0291
0292 bool OrbitNanoAODOutputModule::isFileOpen() const { return nullptr != m_file.get(); }
0293
0294 void OrbitNanoAODOutputModule::openFile(edm::FileBlock const&) {
0295   m_file = std::make_unique<TFile>(m_fileName.c_str(), "RECREATE", "", m_compressionLevel);
0296   edm::Service<edm::JobReport> jr;
0297   cms::Digest branchHash;
0298   m_jrToken = jr->outputFileOpened(m_fileName,
0299                                    m_logicalFileName,
0300                                    std::string(),
0301                                    m_fakeName ? "PoolOutputModule" : "OrbitNanoAODOutputModule",
0302                                    description().moduleLabel(),
0303                                    edm::createGlobalIdentifier(),
0304                                    std::string(),
0305                                    branchHash.digest().toString(),
0306                                    std::vector<std::string>());
0307
0308   if (m_compressionAlgorithm == std::string("ZLIB")) {
0309     m_file->SetCompressionAlgorithm(ROOT::RCompressionSetting::EAlgorithm::kZLIB);
0310   } else if (m_compressionAlgorithm == std::string("LZMA")) {
0311     m_file->SetCompressionAlgorithm(ROOT::RCompressionSetting::EAlgorithm::kLZMA);
0312   } else if (m_compressionAlgorithm == std::string("ZSTD")) {
0313     m_file->SetCompressionAlgorithm(ROOT::RCompressionSetting::EAlgorithm::kZSTD);
0314   } else if (m_compressionAlgorithm == std::string("LZ4")) {
0315     m_file->SetCompressionAlgorithm(ROOT::RCompressionSetting::EAlgorithm::kLZ4);
0316   } else {
0317     throw cms::Exception("Configuration")
0318         << "OrbitNanoAODOutputModule configured with unknown compression algorithm '" << m_compressionAlgorithm << "'\n"
0319         << "Allowed compression algorithms are ZLIB, LZMA, ZSTD, and LZ4\n";
0320   }
0321   /* Setup file structure here */
0322   m_tables.clear();
0323   const auto& keeps = keptProducts();
0324   for (const auto& keep : keeps[edm::InEvent]) {
0325     if (keep.first->className() == "l1ScoutingRun3::OrbitFlatTable")
0326       m_tables.emplace_back(keep.first, keep.second);
0327     else if (keep.first->className() == "std::vector<unsigned int>")
0328       m_selbxs.emplace_back(keep.first, keep.second);
0329     else
0330       throw cms::Exception("Configuration", "OrbitNanoAODOutputModule cannot handle class " + keep.first->className());
0331   }
0332
0333   // create the trees
0334   m_tree = std::make_unique<TTree>("Events", "Events");
0335   m_tree->SetAutoSave(0);
0336   m_tree->SetAutoFlush(0);
0337   m_commonBranches.branch(*m_tree);
0338
0339   m_lumiTree = std::make_unique<TTree>("LuminosityBlocks", "LuminosityBlocks");
0340   m_lumiTree->SetAutoSave(0);
0341   m_commonLumiBranches.branch(*m_lumiTree);
0342
0343   m_runTree = std::make_unique<TTree>("Runs", "Runs");
0344   m_runTree->SetAutoSave(0);
0345   m_commonRunBranches.branch(*m_runTree);
0346
0347   if (m_writeProvenance) {
0348     m_metaDataTree = std::make_unique<TTree>(edm::poolNames::metaDataTreeName().c_str(), "Job metadata");
0349     m_metaDataTree->SetAutoSave(0);
0350     m_parameterSetsTree = std::make_unique<TTree>(edm::poolNames::parameterSetsTreeName().c_str(), "Parameter sets");
0351     m_parameterSetsTree->SetAutoSave(0);
0352   }
0353 }
0354 void OrbitNanoAODOutputModule::reallyCloseFile() {
0355   if (m_writeProvenance) {
0356     int basketSize = 16384;  // fixme configurable?
0357     edm::fillParameterSetBranch(m_parameterSetsTree.get(), basketSize);
0358     edm::fillProcessHistoryBranch(m_metaDataTree.get(), basketSize, m_processHistoryRegistry);
0359     if (m_metaDataTree->GetNbranches() != 0) {
0360       m_metaDataTree->SetEntries(-1);
0361     }
0362     if (m_parameterSetsTree->GetNbranches() != 0) {
0363       m_parameterSetsTree->SetEntries(-1);
0364     }
0365   }
0366   m_file->Write();
0367   m_file->Close();
0368   m_file.reset();
0369   m_tree.release();               // apparently root has ownership
0370   m_lumiTree.release();           //
0371   m_runTree.release();            //
0372   m_metaDataTree.release();       //
0373   m_parameterSetsTree.release();  //
0374   edm::Service<edm::JobReport> jr;
0375   jr->outputFileClosed(m_jrToken);
0376 }
0377
0378 void OrbitNanoAODOutputModule::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0379   edm::ParameterSetDescription desc;
0380
0381   desc.addUntracked<std::string>("fileName");
0382   desc.addUntracked<std::string>("logicalFileName", "");
0383
0384   desc.addUntracked<int>("compressionLevel", 9)->setComment("ROOT compression level of output file.");
0385   desc.addUntracked<std::string>("compressionAlgorithm", "ZLIB")
0386       ->setComment("Algorithm used to compress data in the ROOT output file, allowed values are ZLIB and LZMA");
0387   desc.add<bool>("skipEmptyBXs", false)->setComment("Skip BXs where all input collections are empty");
0388   desc.addUntracked<bool>("saveProvenance", true)
0389       ->setComment("Save process provenance information, e.g. for edmProvDump");
0390   desc.addUntracked<bool>("fakeNameForCrab", false)
0391       ->setComment(
0392           "Change the OutputModule name in the fwk job report to fake PoolOutputModule. This is needed to run on "
0393           "crab "
0394           "(and publish) till crab is fixed");
0395   desc.addUntracked<int>("autoFlush", -10000000)->setComment("Autoflush parameter for ROOT file");
0396
0397   //replace with whatever you want to get from the EDM by default
0398   const std::vector<std::string> keep = {"drop *", "keep l1ScoutingRun3OrbitFlatTable_*Table_*_*"};
0399   edm::one::OutputModule<>::fillDescription(desc, keep);
0400
0401   //Used by Workflow management for their own meta data
0402   edm::ParameterSetDescription dataSet;
0403   dataSet.setAllowAnything();
0404   desc.addUntracked<edm::ParameterSetDescription>("dataset", dataSet)
0405       ->setComment("PSet is only used by Data Operations and not by this module.");
0406
0407   edm::ParameterSetDescription branchSet;
0408   branchSet.setAllowAnything();
0409   desc.add<edm::ParameterSetDescription>("branches", branchSet);
0410
0411   desc.add<edm::InputTag>("selectedBx", edm::InputTag())->setComment("selected Bx (1-3564)");
0412   descriptions.addDefault(desc);
0413 }
0414
0415 DEFINE_FWK_MODULE(OrbitNanoAODOutputModule);