IOPool/Input/src/PoolSource.cc - master - githubmirror/cmssw

IOPool/Input/src/PoolSource.cc
330 строк17 KB
W. David Dagenhart
Reduce the memory usage of the InputFileCatalog
18 май 2026, 23:53
18 май 2026, 23:537e6433c
/*----------------------------------------------------------------------
----------------------------------------------------------------------*/
#include "PoolSource.h"
#include "InputFile.h"
#include "RootPrimaryFileSequence.h"
#include "RootSecondaryFileSequence.h"
#include "DataFormats/Provenance/interface/ProductDescription.h"
#include "DataFormats/Provenance/interface/IndexIntoFile.h"
#include "DataFormats/Provenance/interface/ProductRegistry.h"

#include "FWCore/Framework/interface/EventPrincipal.h"
#include "FWCore/Framework/interface/FileBlock.h"
#include "FWCore/Framework/interface/InputSourceDescription.h"
#include "FWCore/Framework/interface/LuminosityBlockPrincipal.h"
#include "FWCore/Framework/interface/PreallocationConfiguration.h"
#include "FWCore/Framework/interface/SharedResourcesRegistry.h"
#include "FWCore/Framework/interface/SharedResourcesAcquirer.h"
#include "FWCore/Framework/interface/RunPrincipal.h"
#include "FWCore/Framework/interface/ProductResolversFactory.h"
#include "FWCore/Sources/interface/InputSourceRunHelper.h"
#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
#include "FWCore/Utilities/interface/EDMException.h"
#include "FWCore/Utilities/interface/Exception.h"
#include "FWCore/Utilities/interface/InputType.h"

#include <set>

namespace edm {

  class BranchID;
  class LuminosityBlockID;
  class EventID;

  namespace {
    void checkHistoryConsistency(Principal const& primary, Principal const& secondary) {
      ProcessHistory const& ph1 = primary.processHistory();
      ProcessHistory const& ph2 = secondary.processHistory();
      if (ph1 != ph2 && !isAncestor(ph2, ph1)) {
        throw Exception(errors::MismatchedInputFiles, "PoolSource::checkConsistency")
            << "The secondary file is not an ancestor of the primary file\n";
      }
    }
    void checkConsistency(EventPrincipal const& primary, EventPrincipal const& secondary) {
      if (!isSameEvent(primary, secondary)) {
        throw Exception(errors::MismatchedInputFiles, "PoolSource::checkConsistency")
            << primary.id() << " has inconsistent EventAuxiliary data in the primary and secondary file\n";
      }
    }
    void checkConsistency(LuminosityBlockAuxiliary const& primary, LuminosityBlockAuxiliary const& secondary) {
      if (primary.id() != secondary.id()) {
        throw Exception(errors::MismatchedInputFiles, "PoolSource::checkConsistency")
            << primary.id() << " has inconsistent LuminosityBlockAuxiliary data in the primary and secondary file\n";
      }
    }
    void checkConsistency(RunAuxiliary const& primary, RunAuxiliary const& secondary) {
      if (primary.id() != secondary.id()) {
        throw Exception(errors::MismatchedInputFiles, "PoolSource::checkConsistency")
            << primary.id() << " has inconsistent RunAuxiliary data in the primary and secondary file\n";
      }
    }
  }  // namespace

  PoolSource::PoolSource(ParameterSet const& pset, InputSourceDescription const& desc)
      : InputSource(pset, desc),
        rootServiceChecker_(),
        catalog_(pset),
        secondaryCatalog_(pset.getUntrackedParameter<std::vector<std::string> >("secondaryFileNames"),
                          pset.getUntrackedParameter<std::string>("overrideCatalog")),
        secondaryRunPrincipal_(),
        secondaryLumiPrincipal_(),
        secondaryEventPrincipals_(),
        branchIDsToReplace_(),
        nStreams_(desc.allocations_->numberOfStreams()),
        skipBadFiles_(pset.getUntrackedParameter<bool>("skipBadFiles")),
        bypassVersionCheck_(pset.getUntrackedParameter<bool>("bypassVersionCheck")),
        treeMaxVirtualSize_(pset.getUntrackedParameter<int>("treeMaxVirtualSize")),
        productSelectorRules_(pset, "inputCommands", "InputSource"),
        dropDescendants_(pset.getUntrackedParameter<bool>("dropDescendantsOfDroppedBranches")),
        labelRawDataLikeMC_(pset.getUntrackedParameter<bool>("labelRawDataLikeMC")),
        delayReadingEventProducts_(pset.getUntrackedParameter<bool>("delayReadingEventProducts")),
        runHelper_(makeInputSourceRunHelper(pset)),
        resourceSharedWithDelayedReaderPtr_(),
        // Note: primaryFileSequence_ and secondaryFileSequence_ need to be initialized last, because they use data members
        // initialized previously in their own initialization.
        primaryFileSequence_(new RootPrimaryFileSequence(pset, *this, catalog_)),
        secondaryFileSequence_(
            secondaryCatalog_.empty() ? nullptr : new RootSecondaryFileSequence(pset, *this, secondaryCatalog_)) {
    auto resources = SharedResourcesRegistry::instance()->createAcquirerForSourceDelayedReader();
    resourceSharedWithDelayedReaderPtr_ = std::make_unique<SharedResourcesAcquirer>(std::move(resources.first));
    mutexSharedWithDelayedReader_ = resources.second;

    if (secondaryCatalog_.empty() && pset.getUntrackedParameter<bool>("needSecondaryFileNames", false)) {
      throw Exception(errors::Configuration, "PoolSource") << "'secondaryFileNames' must be specified\n";
    }
    if (secondaryFileSequence_) {
      secondaryEventPrincipals_.reserve(nStreams_);
      for (unsigned int index = 0; index < nStreams_; ++index) {
        secondaryEventPrincipals_.emplace_back(new EventPrincipal(secondaryFileSequence_->fileProductRegistry(),
                                                                  edm::productResolversFactory::makePrimary,
                                                                  secondaryFileSequence_->fileBranchIDListHelper(),
                                                                  processConfiguration(),
                                                                  nullptr,
                                                                  index));
      }
      std::array<std::set<BranchID>, NumBranchTypes> idsToReplace;
      ProductRegistry::ProductList const& secondary = secondaryFileSequence_->fileProductRegistry()->productList();
      ProductRegistry::ProductList const& primary = primaryFileSequence_->fileProductRegistry()->productList();
      //this is the registry used by the 'outside' world and only has the primary file information in it at present
      ProductRegistry::ProductList& fullList = productRegistryUpdate().productListUpdator();
      for (auto const& item : secondary) {
        if (item.second.present()) {
          idsToReplace[item.second.branchType()].insert(item.second.branchID());
          //now make sure this is marked as not dropped else the product will not be 'get'table from the Event
          auto itFound = fullList.find(item.first);
          if (itFound != fullList.end()) {
            // If the branch in primary file was dropped, need to initilize the dictionary information
            if (itFound->second.dropped()) {
              itFound->second.initFromDictionary();
            }
            itFound->second.setDropped(false);
          }
        }
      }
      for (auto const& item : primary) {
        if (item.second.present()) {
          idsToReplace[item.second.branchType()].erase(item.second.branchID());
        }
      }
      if (idsToReplace[InEvent].empty() && idsToReplace[InLumi].empty() && idsToReplace[InRun].empty()) {
        secondaryFileSequence_ = nullptr;  // propagate_const<T> has no reset() function
      } else {
        for (int i = InEvent; i < NumBranchTypes; ++i) {
          branchIDsToReplace_[i].reserve(idsToReplace[i].size());
          for (auto const& id : idsToReplace[i]) {
            branchIDsToReplace_[i].push_back(id);
          }
        }
      }
    }
  }

  PoolSource::~PoolSource() {}

  void PoolSource::endJob() {
    if (secondaryFileSequence_)
      secondaryFileSequence_->endJob();
    primaryFileSequence_->endJob();
    InputFile::reportReadBranches();
  }

  std::shared_ptr<FileBlock> PoolSource::readFile_() {
    std::shared_ptr<FileBlock> fb = primaryFileSequence_->readFile_();
    if (secondaryFileSequence_) {
      fb->setNotFastClonable(FileBlock::HasSecondaryFileSequence);
    }
    return fb;
  }

  void PoolSource::closeFile_() { primaryFileSequence_->closeFile(); }

  std::shared_ptr<RunAuxiliary> PoolSource::readRunAuxiliary_() { return primaryFileSequence_->readRunAuxiliary_(); }

  std::shared_ptr<LuminosityBlockAuxiliary> PoolSource::readLuminosityBlockAuxiliary_() {
    return primaryFileSequence_->readLuminosityBlockAuxiliary_();
  }

  void PoolSource::fillProcessBlockHelper_() { primaryFileSequence_->fillProcessBlockHelper_(); }

  bool PoolSource::nextProcessBlock_(ProcessBlockPrincipal& processBlockPrincipal) {
    return primaryFileSequence_->nextProcessBlock_(processBlockPrincipal);
  }

  void PoolSource::readProcessBlock_(ProcessBlockPrincipal& processBlockPrincipal) {
    primaryFileSequence_->readProcessBlock_(processBlockPrincipal);
  }

  void PoolSource::readRun_(RunPrincipal& runPrincipal) {
    bool shouldWeProcessRun = primaryFileSequence_->readRun_(runPrincipal);
    if (secondaryFileSequence_ && shouldWeProcessRun && !branchIDsToReplace_[InRun].empty()) {
      bool found = secondaryFileSequence_->skipToItem(runPrincipal.run(), 0U, 0U);
      if (found) {
        std::shared_ptr<RunAuxiliary> secondaryAuxiliary = secondaryFileSequence_->readRunAuxiliary_();
        checkConsistency(runPrincipal.aux(), *secondaryAuxiliary);
        secondaryRunPrincipal_ = std::make_shared<RunPrincipal>(secondaryFileSequence_->fileProductRegistry(),
                                                                edm::productResolversFactory::makePrimary,
                                                                processConfiguration(),
                                                                nullptr,
                                                                runPrincipal.index());
        secondaryRunPrincipal_->setAux(*secondaryAuxiliary);
        secondaryFileSequence_->readRun_(*secondaryRunPrincipal_);
        checkHistoryConsistency(runPrincipal, *secondaryRunPrincipal_);
        runPrincipal.recombine(*secondaryRunPrincipal_, branchIDsToReplace_[InRun]);
      } else {
        throw Exception(errors::MismatchedInputFiles, "PoolSource::readRun_")
            << " Run " << runPrincipal.run() << " is not found in the secondary input files\n";
      }
    }
  }

  void PoolSource::readLuminosityBlock_(LuminosityBlockPrincipal& lumiPrincipal) {
    bool shouldWeProcessLumi = primaryFileSequence_->readLuminosityBlock_(lumiPrincipal);
    if (secondaryFileSequence_ && shouldWeProcessLumi && !branchIDsToReplace_[InLumi].empty()) {
      bool found = secondaryFileSequence_->skipToItem(lumiPrincipal.run(), lumiPrincipal.luminosityBlock(), 0U);
      if (found) {
        std::shared_ptr<LuminosityBlockAuxiliary> secondaryAuxiliary =
            secondaryFileSequence_->readLuminosityBlockAuxiliary_();
        checkConsistency(lumiPrincipal.aux(), *secondaryAuxiliary);
        secondaryLumiPrincipal_ =
            std::make_shared<LuminosityBlockPrincipal>(secondaryFileSequence_->fileProductRegistry(),
                                                       edm::productResolversFactory::makePrimary,
                                                       processConfiguration(),
                                                       nullptr,
                                                       lumiPrincipal.index());
        secondaryLumiPrincipal_->setAux(*secondaryAuxiliary);
        secondaryFileSequence_->readLuminosityBlock_(*secondaryLumiPrincipal_);
        checkHistoryConsistency(lumiPrincipal, *secondaryLumiPrincipal_);
        lumiPrincipal.recombine(*secondaryLumiPrincipal_, branchIDsToReplace_[InLumi]);
      } else {
        throw Exception(errors::MismatchedInputFiles, "PoolSource::readLuminosityBlock_")
            << " Run " << lumiPrincipal.run() << " LuminosityBlock " << lumiPrincipal.luminosityBlock()
            << " is not found in the secondary input files\n";
      }
    }
  }

  void PoolSource::readEvent_(EventPrincipal& eventPrincipal) {
    bool readAllProducts = not delayReadingEventProducts_;
    bool readEventSucceeded = primaryFileSequence_->readEvent(eventPrincipal, readAllProducts);
    assert(readEventSucceeded);
    if (secondaryFileSequence_ && !branchIDsToReplace_[InEvent].empty()) {
      bool found = secondaryFileSequence_->skipToItem(
          eventPrincipal.run(), eventPrincipal.luminosityBlock(), eventPrincipal.id().event());
      if (found) {
        EventPrincipal& secondaryEventPrincipal = *secondaryEventPrincipals_[eventPrincipal.streamID().value()];
        bool readEventSucceeded = secondaryFileSequence_->readEvent(secondaryEventPrincipal, readAllProducts);
        checkConsistency(eventPrincipal, secondaryEventPrincipal);
        checkHistoryConsistency(eventPrincipal, secondaryEventPrincipal);
        assert(readEventSucceeded);
        eventPrincipal.recombine(secondaryEventPrincipal, branchIDsToReplace_[InEvent]);
        eventPrincipal.mergeProvenanceRetrievers(secondaryEventPrincipal);
        secondaryEventPrincipal.clearPrincipal();
      } else {
        throw Exception(errors::MismatchedInputFiles, "PoolSource::readEvent_")
            << eventPrincipal.id() << " is not found in the secondary input files\n";
      }
    }
    if (readAllProducts) {
      eventPrincipal.readAllFromSourceAndMergeImmediately();
    }
  }

  bool PoolSource::readIt(EventID const& id, EventPrincipal& eventPrincipal, StreamContext& streamContext) {
    bool found = primaryFileSequence_->skipToItem(id.run(), id.luminosityBlock(), id.event());
    if (!found)
      return false;
    EventSourceSentry sentry(*this, streamContext);
    readEvent_(eventPrincipal);
    return true;
  }

  InputSource::ItemTypeInfo PoolSource::getNextItemType() {
    RunNumber_t run = IndexIntoFile::invalidRun;
    LuminosityBlockNumber_t lumi = IndexIntoFile::invalidLumi;
    EventNumber_t event = IndexIntoFile::invalidEvent;
    InputSource::ItemType itemType = primaryFileSequence_->getNextItemType(run, lumi, event);
    if (secondaryFileSequence_ && (ItemType::IsSynchronize != state())) {
      if (itemType == ItemType::IsRun || itemType == ItemType::IsLumi || itemType == ItemType::IsEvent) {
        if (!secondaryFileSequence_->containedInCurrentFile(run, lumi, event)) {
          return ItemTypeInfo::isSynchronize();
        }
      }
    }
    return InputSource::ItemTypeInfo(runHelper_->nextItemType(state(), itemType, run, lumi, event));
  }

  std::pair<SharedResourcesAcquirer*, std::recursive_mutex*> PoolSource::resourceSharedWithDelayedReader_() {
    return std::make_pair(resourceSharedWithDelayedReaderPtr_.get(), mutexSharedWithDelayedReader_.get());
  }

  // Rewind to before the first event that was read.
  void PoolSource::rewind_() { primaryFileSequence_->rewind_(); }

  // Advance "offset" events.  Offset can be positive or negative (or zero).
  void PoolSource::skip(int offset) { primaryFileSequence_->skipEvents(offset); }

  bool PoolSource::goToEvent_(EventID const& eventID) { return primaryFileSequence_->goToEvent(eventID); }

  void PoolSource::fillDescriptions(ConfigurationDescriptions& descriptions) {
    ParameterSetDescription desc;

    std::vector<std::string> defaultStrings;
    desc.setComment("Reads EDM/Root files.");
    InputFileCatalog::fillDescription(desc);
    desc.addUntracked<std::vector<std::string> >("secondaryFileNames", defaultStrings)
        ->setComment("Names of secondary files to be processed.");
    desc.addUntracked<bool>("needSecondaryFileNames", false)
        ->setComment("If True, 'secondaryFileNames' must be specified and be non-empty.");
    desc.addUntracked<bool>("skipBadFiles", false)
        ->setComment(
            "True:  Ignore any missing or unopenable input file.\n"
            "False: Throw exception if missing or unopenable input file.");
    desc.addUntracked<bool>("bypassVersionCheck", false)
        ->setComment(
            "True:  Bypass release version check.\n"
            "False: Throw exception if reading file in a release prior to the release in which the file was written.");
    desc.addUntracked<int>("treeMaxVirtualSize", -1)
        ->setComment("Size of ROOT TTree TBasket cache. Affects performance.");
    desc.addUntracked<bool>("dropDescendantsOfDroppedBranches", true)
        ->setComment("If True, also drop on input any descendent of any branch dropped on input.");
    desc.addUntracked<bool>("labelRawDataLikeMC", true)
        ->setComment("If True: replace module label for raw data to match MC. Also use 'LHC' as process.");
    desc.addUntracked<bool>("delayReadingEventProducts", true)
        ->setComment(
            "If True: do not read a data product from the file until it is requested. If False: all event data "
            "products are read upfront.");
    ProductSelectorRules::fillDescription(desc, "inputCommands");
    InputSource::fillDescription(desc);
    RootPrimaryFileSequence::fillDescription(desc);
    InputSourceRunHelperBase::fillDescription(desc);

    descriptions.add("source", desc);
  }

  bool PoolSource::randomAccess_() const { return true; }

  ProcessingController::ForwardState PoolSource::forwardState_() const { return primaryFileSequence_->forwardState(); }

  ProcessingController::ReverseState PoolSource::reverseState_() const { return primaryFileSequence_->reverseState(); }
}  // namespace edm