Support skipping invalid timeframes across parent files

Author: autumn-mck

Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx

@@ -275,7 +275,8 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
         auto skippedTimeframes = ++totalInvalidReadSkipped;
         LOGP(error, "Invalid AOD read for table {}: fileCounter {}, timeFrame {}. Skipping timeframe (skipped timeframes: {}). Reason: {}",
              concrete.origin.as<std::string>(), fcnt, ntf, skippedTimeframes, describeException(e));
-        didir->markTimeFrameSkipped(header::DataHeader(concrete.description, concrete.origin, concrete.subSpec), ntf);
+        clean_all_runtime_errors();
+        didir->finishTimeFrame(true);
         arrowContext.clear();
         messageContext.discard();
         stringContext.clear();
@@ -354,6 +355,9 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
         auto dh = header::DataHeader(concrete.description, concrete.origin, concrete.subSpec);
         bool wasAOD = std::ranges::any_of(route.matcher.metadata, [](ConfigParamSpec const& p) { return p.name.starts_with("aod-origin-replaced"); });
 
+        if (currentState == TFReaderState::READ_FIRST_TABLE || currentState == TFReaderState::READ_FIRST_TABLE_FROM_NEXT_FILE) {
+          didir->beginTimeFrame();
+        }
         try {
           if (!didir->readTree(outputs, dh, fcnt, ntf, totalSizeCompressed, totalSizeUncompressed, wasAOD)) {
             return TFReaderState::TRY_NEXT_FILE;
@@ -388,6 +392,7 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
             }
             break;
           case TFReaderState::TRY_NEXT_FILE:
+            didir->finishTimeFrame();
             fcnt += device.maxInputTimeslices;
             if (didir->atEnd(fcnt)) {
               LOGP(info, "No input files left to read for reader {}!", device.inputTimesliceId);
@@ -407,6 +412,7 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
             break;
         }
       }
+      didir->finishTimeFrame();
       int64_t stopSize = totalSizeCompressed;
       int64_t bytesDelta = stopSize - startSize;
       int64_t stopTime = uv_hrtime();

Framework/AnalysisSupport/src/DataInputDirector.cxx

@@ -161,10 +161,13 @@ bool DataInputDescriptor::setFile(int counter, int wantedParentLevel, std::strin
   if (tfile == nullptr) {
     tfile = TFile::Open(filename.c_str());
   }
-  mCurrentFilesystem = std::make_shared<TFileFileSystem>(tfile, 50 * 1024 * 1024, mFactory, !externalFile);
-  if (!mCurrentFilesystem.get()) {
+  if (!tfile || tfile->IsZombie()) {
+    if (tfile && !externalFile) {
+      delete tfile;
+    }
     throw std::runtime_error(fmt::format("Couldn't open file \"{}\"!", filename));
   }
+  mCurrentFilesystem = std::make_shared<TFileFileSystem>(tfile, 50 * 1024 * 1024, mFactory, !externalFile);
   rootFS = std::dynamic_pointer_cast<TFileFileSystem>(mCurrentFilesystem);
   printFileOpening();
 
@@ -254,6 +257,10 @@ std::pair<std::shared_ptr<DataInputDescriptor>, int> DataInputDescriptor::naviga
   if (!setFile(counter, wantedParentLevel, wantedOrigin)) {
     return {nullptr, -1};
   }
+  if (numTF < 0 || numTF >= mfilenames[counter].numberOfTimeFrames) {
+    return {nullptr, -1};
+  }
+  recordTimeFrameRead(counter, numTF);
   auto folderName = fmt::format("DF_{}", mfilenames[counter].listOfTimeFrameNumbers[numTF]);
   auto parentFile = getParentFile(counter, numTF, "", wantedParentLevel, wantedOrigin);
   if (parentFile == nullptr) {
@@ -283,18 +290,65 @@ arrow::dataset::FileSource DataInputDescriptor::getFileFolder(int counter, int n
     return {};
   }
 
+  recordTimeFrameRead(counter, numTF);
   mfilenames[counter].alreadyRead[numTF] = true;
 
   return {fmt::format("DF_{}", mfilenames[counter].listOfTimeFrameNumbers[numTF]), mCurrentFilesystem};
 }
 
-uint64_t DataInputDescriptor::markTimeFrameSkipped(int numTF)
+void DataInputDescriptor::recordTimeFrameRead(int counter, int numTF)
 {
-  if (mCurrentFileID >= 0 && numTF >= 0 && numTF < mfilenames[mCurrentFileID].numberOfTimeFrames) {
-    mfilenames[mCurrentFileID].alreadyRead[numTF] = false;
-    return ++mfilenames[mCurrentFileID].invalidReadSkipped;
+  auto read = std::pair{counter, numTF};
+  if (std::find(mTimeFrameReads.begin(), mTimeFrameReads.end(), read) == mTimeFrameReads.end()) {
+    mTimeFrameReads.push_back(read);
+  }
+}
+
+void DataInputDescriptor::beginTimeFrame()
+{
+  mTimeFrameReads.clear();
+  mTimeFrameActive = true;
+  if (mParentFile) {
+    mParentFile->beginTimeFrame();
+  }
+}
+
+void DataInputDescriptor::finishTimeFrame(bool skipped)
+{
+  mTimeFrameActive = false;
+  if (skipped) {
+    for (auto [file, df] : mTimeFrameReads) {
+      mfilenames[file].alreadyRead[df] = false;
+      ++mfilenames[file].invalidReadSkipped;
+    }
+  }
+  mTimeFrameReads.clear();
+  if (mParentFile) {
+    mParentFile->finishTimeFrame(skipped);
+  }
+  for (auto& parent : mRetainedParents) {
+    parent->finishTimeFrame(skipped);
+    parent->closeInputFile();
+  }
+  mRetainedParents.clear();
+  for (auto& [counter, info] : mPendingFileStatistics) {
+    reportFileStatistics(counter, std::move(info));
+  }
+  mPendingFileStatistics.clear();
+}
+
+void DataInputDescriptor::releaseParentFile()
+{
+  if (!mParentFile) {
+    return;
+  }
+  if (mTimeFrameActive) {
+    // Keep both the read bookkeeping and the open file for final statistics.
+    mRetainedParents.push_back(std::move(mParentFile));
+  } else {
+    mParentFile->closeInputFile();
+    mParentFile.reset();
   }
-  return 0;
 }
 
 std::shared_ptr<DataInputDescriptor> DataInputDescriptor::getParentFile(int counter, int numTF, std::string treename, int wantedParentLevel, std::string_view wantedOrigin)
@@ -309,8 +363,7 @@ std::shared_ptr<DataInputDescriptor> DataInputDescriptor::getParentFile(int coun
   // The current DF is not found in the parent map (this should not happen and is a fatal error)
   auto rootFS = std::dynamic_pointer_cast<TFileFileSystem>(mCurrentFilesystem);
   if (!parentFileName) {
-    throw std::runtime_error(fmt::format(R"(parent file map exists but does not contain the current DF "{}" in file "{}")", folderName.c_str(), rootFS->GetFile()->GetName()));
-    return nullptr;
+    throw InvalidAODReadError(fmt::format(R"(parent file map exists but does not contain the current DF "{}" in file "{}")", folderName.c_str(), rootFS->GetFile()->GetName()));
   }
 
   if (mParentFile) {
@@ -319,21 +372,28 @@ std::shared_ptr<DataInputDescriptor> DataInputDescriptor::getParentFile(int coun
     if (parentFileName->GetString().CompareTo(parentRootFS->GetFile()->GetName()) == 0) {
       return mParentFile;
     } else {
-      mParentFile->closeInputFile();
-      mParentFile.reset();
+      releaseParentFile();
     }
   }
 
   if (mLevel == mContext.allowedParentLevel) {
-    throw std::runtime_error(fmt::format(R"(while looking for tree "{}", the parent file was requested but we are already at level {} of maximal allowed level {} for DF "{}" in file "{}")", treename.c_str(), mLevel, mContext.allowedParentLevel, folderName.c_str(),
-                                         rootFS->GetFile()->GetName()));
+    throw InvalidAODReadError(fmt::format(R"(while looking for tree "{}", the parent file was requested but we are already at level {} of maximal allowed level {} for DF "{}" in file "{}")", treename.c_str(), mLevel, mContext.allowedParentLevel, folderName.c_str(),
+                                          rootFS->GetFile()->GetName()));
   }
 
   LOGP(info, "Opening parent file {} for DF {}", parentFileName->GetString().Data(), folderName.c_str());
   mParentFile = std::make_shared<DataInputDescriptor>(mAlienSupport, mLevel + 1, mContext);
   mParentFile->mdefaultFilenamesPtr.emplace_back(makeFileNameHolder(parentFileName->GetString().Data()));
   mParentFile->fillInputfiles();
-  mParentFile->setFile(0, wantedParentLevel, wantedOrigin);
+  try {
+    mParentFile->setFile(0, wantedParentLevel, wantedOrigin);
+    if (mTimeFrameActive) {
+      mParentFile->beginTimeFrame();
+    }
+  } catch (...) {
+    mParentFile.reset();
+    std::throw_with_nested(InvalidAODReadError(fmt::format("Unable to open parent file \"{}\" for {}", parentFileName->GetString().Data(), folderName)));
+  }
   return mParentFile;
 }
 
@@ -373,15 +433,26 @@ void DataInputDescriptor::printFileStatistics()
   }
   auto rootFS = std::dynamic_pointer_cast<TFileFileSystem>(mCurrentFilesystem);
   auto f = dynamic_cast<TFile*>(rootFS->GetFile());
-  std::string monitoringInfo(fmt::format("lfn={},size={},total_df={},read_df={},skipped_df={},read_bytes={},read_calls={},io_time={:.1f},wait_time={:.1f},level={}", f->GetName(),
-                                         f->GetSize(), getTimeFramesInFile(mCurrentFileID), getReadTimeFramesInFile(mCurrentFileID), mfilenames.at(mCurrentFileID).invalidReadSkipped, f->GetBytesRead(), f->GetReadCalls(),
+  std::string monitoringInfo(fmt::format("lfn={},size={},total_df={},read_bytes={},read_calls={},io_time={:.1f},wait_time={:.1f},level={}", f->GetName(),
+                                         f->GetSize(), getTimeFramesInFile(mCurrentFileID), f->GetBytesRead(), f->GetReadCalls(),
                                          ((float)mIOTime / 1e9), ((float)wait_time / 1e9), mLevel));
 #if __has_include(<TJAlienFile.h>)
   auto alienFile = dynamic_cast<TJAlienFile*>(f);
   if (alienFile) {
     monitoringInfo += fmt::format(",se={},open_time={:.1f}", alienFile->GetSE(), alienFile->GetElapsed());
   }
 #endif
+  if (mTimeFrameActive) {
+    // Snapshot I/O statistics now, but publish DF counts after commit or rollback.
+    mPendingFileStatistics.emplace_back(mCurrentFileID, std::move(monitoringInfo));
+  } else {
+    reportFileStatistics(mCurrentFileID, std::move(monitoringInfo));
+  }
+}
+
+void DataInputDescriptor::reportFileStatistics(int counter, std::string monitoringInfo)
+{
+  monitoringInfo += fmt::format(",read_df={},skipped_df={}", getReadTimeFramesInFile(counter), mfilenames.at(counter).invalidReadSkipped);
   if (mContext.monitoring) {
     mContext.monitoring->send(o2::monitoring::Metric{monitoringInfo, "aod-file-read-info"}.addTag(o2::monitoring::tags::Key::Subsystem, o2::monitoring::tags::Value::DPL));
   }
@@ -391,10 +462,7 @@ void DataInputDescriptor::printFileStatistics()
 void DataInputDescriptor::closeInputFile()
 {
   if (mCurrentFilesystem.get()) {
-    if (mParentFile) {
-      mParentFile->closeInputFile();
-      mParentFile.reset();
-    }
+    releaseParentFile();
 
     delete mParentFileMap;
     mParentFileMap = nullptr;
@@ -492,7 +560,7 @@ struct CalculateDelta {
 };
 
 bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh, int counter, int numTF, std::string treename, size_t& totalSizeCompressed, size_t& totalSizeUncompressed)
-{
+try {
   CalculateDelta t(mIOTime);
   std::string wantedOrigin = dh.dataOrigin.as<std::string>();
   int wantedLevel = mContext.levelForOrigin(wantedOrigin);
@@ -501,13 +569,17 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
   // attempting to read from this level.
   if (wantedLevel != -1 && mLevel < wantedLevel) {
     auto [parentFile, parentNumTF] = navigateToLevel(counter, numTF, wantedLevel, wantedOrigin);
+    if (counter >= getNumberInputfiles() || numTF < 0 || numTF >= mfilenames[counter].numberOfTimeFrames) {
+      t.deactivate();
+      return false;
+    }
     if (parentFile == nullptr) {
       auto rootFS = std::dynamic_pointer_cast<TFileFileSystem>(mCurrentFilesystem);
-      throw std::runtime_error(fmt::format(R"(No parent file found for "{}" while looking for level {} in "{}")", treename, wantedLevel, rootFS->GetFile()->GetName()));
+      throw InvalidAODReadError(fmt::format(R"(No parent file found for "{}" while looking for level {} in "{}")", treename, wantedLevel, rootFS->GetFile()->GetName()));
     }
     if (parentNumTF == -1) {
       auto parentRootFS = std::dynamic_pointer_cast<TFileFileSystem>(parentFile->mCurrentFilesystem);
-      throw std::runtime_error(fmt::format(R"(DF not found in parent file "{}")", parentRootFS->GetFile()->GetName()));
+      throw InvalidAODReadError(fmt::format(R"(DF not found in parent file "{}")", parentRootFS->GetFile()->GetName()));
     }
     t.deactivate();
     return parentFile->readTree(outputs, dh, 0, parentNumTF, treename, totalSizeCompressed, totalSizeUncompressed);
@@ -546,12 +618,7 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
   if (!format) {
     t.deactivate();
     LOGP(debug, "Could not find tree {}. Trying in parent file.", fullpath.path());
-    std::shared_ptr<DataInputDescriptor> parentFile;
-    try {
-      parentFile = getParentFile(counter, numTF, treename, wantedLevel, wantedOrigin);
-    } catch (...) {
-      std::throw_with_nested(InvalidAODReadError(fmt::format("Unable to resolve parent file for tree {}", treename)));
-    }
+    auto parentFile = getParentFile(counter, numTF, treename, wantedLevel, wantedOrigin);
     if (parentFile == nullptr) {
       auto rootFS = std::dynamic_pointer_cast<TFileFileSystem>(mCurrentFilesystem);
       throw std::runtime_error(fmt::format(R"(Couldn't get TTree "{}" from "{}". Please check https://aliceo2group.github.io/analysis-framework/docs/troubleshooting/#tree-not-found for more information.)", fullpath.path(), rootFS->GetFile()->GetName()));
@@ -599,6 +666,12 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
   }
 
   return true;
+} catch (InvalidAODReadError const&) {
+  auto rootFS = std::dynamic_pointer_cast<TFileFileSystem>(mCurrentFilesystem);
+  auto filename = rootFS ? std::string(rootFS->GetFile()->GetName()) : mfilenames.at(counter).fileName;
+  auto const& dfs = mfilenames.at(counter).listOfTimeFrameNumbers;
+  auto df = numTF >= 0 && static_cast<size_t>(numTF) < dfs.size() ? fmt::format("DF_{}", dfs[numTF]) : fmt::format("timeframe index {}", numTF);
+  std::throw_with_nested(InvalidAODReadError(fmt::format("Reading tree {} in {} from file \"{}\" at parent level {}", treename, df, filename, mLevel)));
 }
 
 DataInputDirector::DataInputDirector(std::vector<std::string> inputFiles, DataInputDirectorContext&& context)
@@ -888,13 +961,20 @@ arrow::dataset::FileSource DataInputDirector::getFileFolder(header::DataHeader d
   return didesc->getFileFolder(counter, numTF, wantedLevel, origin);
 }
 
-void DataInputDirector::markTimeFrameSkipped(header::DataHeader dh, int numTF)
+void DataInputDirector::beginTimeFrame()
 {
-  auto didesc = getDataInputDescriptor(dh);
-  if (!didesc) {
-    didesc = mdefaultDataInputDescriptor.get();
+  mdefaultDataInputDescriptor->beginTimeFrame();
+  for (auto& descriptor : mdataInputDescriptors) {
+    descriptor.beginTimeFrame();
+  }
+}
+
+void DataInputDirector::finishTimeFrame(bool skipped)
+{
+  mdefaultDataInputDescriptor->finishTimeFrame(skipped);
+  for (auto& descriptor : mdataInputDescriptors) {
+    descriptor.finishTimeFrame(skipped);
   }
-  didesc->markTimeFrameSkipped(numTF);
 }
 
 int DataInputDirector::getTimeFramesInFile(header::DataHeader dh, int counter)
@@ -944,6 +1024,7 @@ bool DataInputDirector::readTree(DataAllocator& outputs, header::DataHeader dh,
 
 void DataInputDirector::closeInputFiles()
 {
+  finishTimeFrame();
   mdefaultDataInputDescriptor->closeInputFile();
   for (auto& didesc : mdataInputDescriptors) {
     didesc.closeInputFile();

Framework/AnalysisSupport/src/DataInputDirector.h

@@ -107,7 +107,10 @@ class DataInputDescriptor
 
   uint64_t getTimeFrameNumber(int counter, int numTF, int wantedParentLevel, std::string_view wantedOrigin);
   arrow::dataset::FileSource getFileFolder(int counter, int numTF, int wantedParentLevel, std::string_view wantedOrigin);
-  uint64_t markTimeFrameSkipped(int numTF);
+  // Start tracking the file/DF pairs participating in one output timeframe.
+  void beginTimeFrame();
+  // Commit or discard a timeframe, then close retained parents and report statistics.
+  void finishTimeFrame(bool skipped = false);
   // Open the current file to populate the parent map, then return the parent descriptor and
   // the TF index within it that corresponds to numTF at this level. Returns {nullptr, -1} on failure.
   std::pair<std::shared_ptr<DataInputDescriptor>, int> navigateToLevel(int counter, int numTF, int wantedParentLevel, std::string_view wantedOrigin);
@@ -123,6 +126,10 @@ class DataInputDescriptor
   bool isAlienSupportOn() { return mAlienSupport; }
 
  private:
+  void recordTimeFrameRead(int counter, int numTF);
+  void releaseParentFile();
+  void reportFileStatistics(int counter, std::string monitoringInfo);
+
   o2::framework::RootObjectReadingFactory mFactory;
   std::string minputfilesFile;
   std::string* minputfilesFilePtr = nullptr;
@@ -132,11 +139,16 @@ class DataInputDescriptor
   std::vector<FileNameHolder> mdefaultFilenamesPtr;
   std::shared_ptr<arrow::fs::FileSystem> mCurrentFilesystem;
   int mCurrentFileID = -1;
+  // File-local DF indices visited during the current output timeframe.
+  std::vector<std::pair<int, int>> mTimeFrameReads;
   bool mAlienSupport = false;
 
   DataInputDirectorContext& mContext;
   TMap* mParentFileMap = nullptr;
   std::shared_ptr<DataInputDescriptor> mParentFile = nullptr;
+  std::vector<std::shared_ptr<DataInputDescriptor>> mRetainedParents;
+  bool mTimeFrameActive = false;
+  std::vector<std::pair<int, std::string>> mPendingFileStatistics;
   int mLevel = 0; // level of parent files
 
   int mtotalNumberTimeFrames = 0;
@@ -173,7 +185,10 @@ class DataInputDirector
   bool readTree(DataAllocator& outputs, header::DataHeader dh, int counter, int numTF, size_t& totalSizeCompressed, size_t& totalSizeUncompressed, bool wasAOD);
   uint64_t getTimeFrameNumber(header::DataHeader dh, int counter, int numTF);
   arrow::dataset::FileSource getFileFolder(header::DataHeader dh, int counter, int numTF);
-  void markTimeFrameSkipped(header::DataHeader dh, int numTF);
+  // Start tracking the file/DF pairs participating in one output timeframe.
+  void beginTimeFrame();
+  // Commit or discard a timeframe, then close retained parents and report statistics.
+  void finishTimeFrame(bool skipped = false);
   int getTimeFramesInFile(header::DataHeader dh, int counter);
 
   uint64_t getTotalSizeCompressed();