Skip invalid timeframes during ROOT input

Handle corrupt reads as recoverable and discard the affected timeframe when DPL_AOD_READER_SKIP_INVALID is enabled.

Author: autumn-mck

Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx

@@ -11,6 +11,7 @@
 
 #include "AODJAlienReaderHelpers.h"
 #include <charconv>
+#include <cstdlib>
 #include <memory>
 #include <ranges>
 #include <vector>
@@ -19,13 +20,16 @@
 #include "Framework/DataProcessingStats.h"
 #include "Framework/RootArrowFilesystem.h"
 #include "Framework/AlgorithmSpec.h"
+#include "Framework/ArrowContext.h"
 #include "Framework/ConfigParamRegistry.h"
 #include "Framework/ControlService.h"
 #include "Framework/CallbackService.h"
 #include "Framework/EndOfStreamContext.h"
 #include "Framework/DeviceSpec.h"
 #include "Framework/RawDeviceService.h"
 #include "Framework/DataSpecUtils.h"
+#include "Framework/MessageContext.h"
+#include "Framework/StringContext.h"
 #include "Framework/ConfigContext.h"
 #include "DataInputDirector.h"
 #include "Framework/SourceInfoHeader.h"
@@ -200,7 +204,7 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
                            numTF,
                            watchdog,
                            maxRate,
-                           didir, reportTFN, reportTFFileName, level](Monitoring& monitoring, DataAllocator& outputs, ControlService& control, DeviceSpec const& device, DataProcessingStats& dpstats) {
+                           didir, reportTFN, reportTFFileName, level](Monitoring& monitoring, DataAllocator& outputs, ControlService& control, DeviceSpec const& device, DataProcessingStats& dpstats, ArrowContext& arrowContext, MessageContext& messageContext, StringContext& stringContext) {
       // Each parallel reader device.inputTimesliceId reads the files fileCounter*device.maxInputTimeslices+device.inputTimesliceId
       // the TF to read is numTF
       assert(device.inputTimesliceId < device.maxInputTimeslices);
@@ -218,6 +222,8 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
       static size_t totalSizeUncompressed = 0;
       static size_t totalSizeCompressed = 0;
       static uint64_t totalDFSent = 0;
+      static uint64_t totalInvalidReadSkipped = 0;
+      static bool skipInvalidReads = getenv("DPL_AOD_READER_SKIP_INVALID") && atoi(getenv("DPL_AOD_READER_SKIP_INVALID"));
 
       // check if RuntimeLimit is reached
       if (!watchdog->update()) {
@@ -232,6 +238,17 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
 
       int64_t startTime = uv_hrtime();
       int64_t startSize = totalSizeCompressed;
+      auto skipInvalidRead = [&](o2::header::DataOrigin const& origin, InvalidAODReadError const& e) {
+        auto skippedTimeframes = ++totalInvalidReadSkipped;
+        LOGP(error, "Invalid AOD read for table {}: fileCounter {}, timeFrame {}. Skipping timeframe (skipped timeframes: {}). Reason: {}",
+             origin.as<std::string>(), fcnt, ntf, skippedTimeframes, e.what());
+        arrowContext.clear();
+        messageContext.discard();
+        stringContext.clear();
+        monitoring.send(Metric{skippedTimeframes, "aod-invalid-read-skipped-timeframes"}.addTag(Key::Subsystem, monitoring::tags::Value::DPL));
+        *fileCounter = (fcnt - device.inputTimesliceId) / device.maxInputTimeslices;
+        *numTF = ntf;
+      };
       for (auto& route : requestedTables) {
         if ((device.inputTimesliceId % route.maxTimeslices) != route.timeslice) {
           continue;
@@ -242,25 +259,44 @@ 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 (!didir->readTree(outputs, dh, fcnt, ntf, totalSizeCompressed, totalSizeUncompressed, wasAOD)) {
-          if (first) {
-            // check if there is a next file to read
-            fcnt += device.maxInputTimeslices;
-            if (didir->atEnd(fcnt)) {
-              LOGP(info, "No input files left to read for reader {}!", device.inputTimesliceId);
-              didir->closeInputFiles();
-              monitoring.flushBuffer();
-              control.endOfStream();
-              control.readyToQuit(QuitRequest::Me);
-              return;
-            }
-            // get first folder of next file
-            ntf = 0;
-            if (!didir->readTree(outputs, dh, fcnt, ntf, totalSizeCompressed, totalSizeUncompressed, wasAOD)) {
-              LOGP(fatal, "Can not retrieve tree for table {}: fileCounter {}, timeFrame {}", concrete.origin.as<std::string>(), fcnt, ntf);
-              throw std::runtime_error("Processing is stopped!");
+        bool treeRead = false;
+        try {
+          treeRead = didir->readTree(outputs, dh, fcnt, ntf, totalSizeCompressed, totalSizeUncompressed, wasAOD);
+        } catch (InvalidAODReadError const& e) {
+          if (!skipInvalidReads) {
+            throw;
+          }
+          skipInvalidRead(concrete.origin, e);
+          return;
+        }
+
+        if (!treeRead) {
+          if (!first) {
+            LOGP(fatal, "Can not retrieve tree for table {}: fileCounter {}, timeFrame {}", concrete.origin.as<std::string>(), fcnt, ntf);
+            throw std::runtime_error("Processing is stopped!");
+          }
+          // check if there is a next file to read
+          fcnt += device.maxInputTimeslices;
+          if (didir->atEnd(fcnt)) {
+            LOGP(info, "No input files left to read for reader {}!", device.inputTimesliceId);
+            didir->closeInputFiles();
+            monitoring.flushBuffer();
+            control.endOfStream();
+            control.readyToQuit(QuitRequest::Me);
+            return;
+          }
+          // get first folder of next file
+          ntf = 0;
+          try {
+            treeRead = didir->readTree(outputs, dh, fcnt, ntf, totalSizeCompressed, totalSizeUncompressed, wasAOD);
+          } catch (InvalidAODReadError const& e) {
+            if (!skipInvalidReads) {
+              throw;
             }
-          } else {
+            skipInvalidRead(concrete.origin, e);
+            return;
+          }
+          if (!treeRead) {
             LOGP(fatal, "Can not retrieve tree for table {}: fileCounter {}, timeFrame {}", concrete.origin.as<std::string>(), fcnt, ntf);
             throw std::runtime_error("Processing is stopped!");
           }

Framework/AnalysisSupport/src/DataInputDirector.cxx

@@ -34,6 +34,7 @@
 #include <arrow/dataset/file_base.h>
 #include <arrow/dataset/dataset.h>
 #include <uv.h>
+#include <exception>
 #include <memory>
 
 #if __has_include(<TJAlienFile.h>)
@@ -536,18 +537,25 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
   if (!format) {
     t.deactivate();
     LOGP(debug, "Could not find tree {}. Trying in parent file.", fullpath.path());
-    auto parentFile = getParentFile(counter, numTF, treename, wantedLevel, wantedOrigin);
-    if (parentFile != nullptr) {
-      int parentNumTF = parentFile->findDFNumber(0, folder.path());
-      if (parentNumTF == -1) {
-        auto parentRootFS = std::dynamic_pointer_cast<TFileFileSystem>(parentFile->mCurrentFilesystem);
-        throw std::runtime_error(fmt::format(R"(DF {} listed in parent file map but not found in the corresponding file "{}")", folder.path(), parentRootFS->GetFile()->GetName()));
-      }
-      // first argument is 0 as the parent file object contains only 1 file
-      return parentFile->readTree(outputs, dh, 0, parentNumTF, treename, totalSizeCompressed, totalSizeUncompressed);
+    std::shared_ptr<DataInputDescriptor> parentFile;
+    try {
+      parentFile = getParentFile(counter, numTF, treename, wantedLevel, wantedOrigin);
+    } catch (std::exception const& e) {
+      throw InvalidAODReadError(fmt::format("Unable to resolve parent file for tree {}: {}", treename, e.what()));
+    } catch (...) {
+      throw InvalidAODReadError(fmt::format("Unable to resolve parent file for tree {}", treename));
+    }
+    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()));
     }
-    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()));
+    int parentNumTF = parentFile->findDFNumber(0, folder.path());
+    if (parentNumTF == -1) {
+      auto parentRootFS = std::dynamic_pointer_cast<TFileFileSystem>(parentFile->mCurrentFilesystem);
+      throw InvalidAODReadError(fmt::format(R"(DF {} listed in parent file map but not found in the corresponding file "{}")", folder.path(), parentRootFS->GetFile()->GetName()));
+    }
+    // first argument is 0 as the parent file object contains only 1 file
+    return parentFile->readTree(outputs, dh, 0, parentNumTF, treename, totalSizeCompressed, totalSizeUncompressed);
   }
 
   auto schemaOpt = format->Inspect(fullpath);
@@ -573,7 +581,23 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
   //// add branches to read
   //// fill the table
   f2b->setLabel(treename.c_str());
-  f2b->fill(datasetSchema, format);
+  try {
+    f2b->fill(datasetSchema, format);
+  } catch (std::exception const& e) {
+    f2b.discard();
+    throw InvalidAODReadError(fmt::format("Unable to read tree {}: {}", treename, e.what()));
+  } catch (...) {
+    f2b.discard();
+    throw InvalidAODReadError(fmt::format("Unable to read tree {}", treename));
+  }
+
+  try {
+    f2b.release();
+  } catch (std::exception const& e) {
+    throw InvalidAODReadError(fmt::format("Unable to finalize tree {}: {}", treename, e.what()));
+  } catch (...) {
+    throw InvalidAODReadError(fmt::format("Unable to finalize tree {}", treename));
+  }
 
   return true;
 }

Framework/AnalysisSupport/src/DataInputDirector.h

@@ -21,6 +21,7 @@
 #include <arrow/dataset/dataset.h>
 
 #include <regex>
+#include <stdexcept>
 #include <vector>
 #include "rapidjson/fwd.h"
 
@@ -32,6 +33,12 @@ class Monitoring;
 namespace o2::framework
 {
 
+class InvalidAODReadError : public std::runtime_error
+{
+ public:
+  using std::runtime_error::runtime_error;
+};
+
 struct FileNameHolder {
   std::string fileName;
   int numberOfTimeFrames = 0;

Framework/Core/include/Framework/MessageContext.h

@@ -466,6 +466,11 @@ class MessageContext
   /// discarded.
   void clear();
 
+  /// Discard pending output messages without asserting that they were sent. This
+  /// is intended for exception teardown paths where normal post-processing will
+  /// not run.
+  void discard();
+
   FairMQDeviceProxy& proxy()
   {
     return mProxy;

Framework/Core/src/MessageContext.cxx

@@ -107,6 +107,13 @@ void MessageContext::clear()
   mMessages.clear();
 }
 
+void MessageContext::discard()
+{
+  mDidDispatch = false;
+  mScheduledMessages.clear();
+  mMessages.clear();
+}
+
 int64_t MessageContext::addToCache(std::unique_ptr<fair::mq::Message>& toCache)
 {
   auto&& cached = toCache->GetTransport()->CreateMessage();