DPL Analysis: write HistogramRegistry incrementally

ktf · ktf · commit daa8075475e5 · 2025-09-03T13:14:26.000+02:00
This should reduce big spikes at the end of the processing when a large
HistogramRegistry is serialised.
diff --git a/Framework/AnalysisSupport/src/AODWriterHelpers.cxx b/Framework/AnalysisSupport/src/AODWriterHelpers.cxx
@@ -24,12 +24,14 @@
 #include "Framework/Monitoring.h"
 
 #include <Monitoring/Monitoring.h>
+#include <TDirectory.h>
 #include <TFile.h>
 #include <TFile.h>
 #include <TTree.h>
 #include <TMap.h>
 #include <TObjString.h>
 #include <arrow/table.h>
+#include <chrono>
 
 namespace o2::framework::writers
 {
@@ -46,6 +48,7 @@ struct InputObjectRoute {
 struct InputObject {
   TClass* kind = nullptr;
   void* obj = nullptr;
+  std::string container;
   std::string name;
   int count = -1;
 };
@@ -312,6 +315,13 @@ AlgorithmSpec AODWriterHelpers::getOutputObjHistWriter(ConfigContext const& ctx)
         obj.obj = tm.ReadObjectAny(obj.kind);
         auto* named = static_cast<TNamed*>(obj.obj);
         obj.name = named->GetName();
+        // If we have a folder, we assume the first element of the path
+        // to be the name of the registry.
+        if (sourceType == HistogramRegistrySource) {
+          obj.container = objh->containerName;
+        } else {
+          obj.container = obj.name;
+        }
         auto hpos = std::find_if(tskmap.begin(), tskmap.end(), [&](auto&& x) { return x.id == hash; });
         if (hpos == tskmap.end()) {
           LOG(error) << "No task found for hash " << hash;
@@ -324,8 +334,8 @@ AlgorithmSpec AODWriterHelpers::getOutputObjHistWriter(ConfigContext const& ctx)
           return;
         }
         auto objects = opos->bindings;
-        if (std::find(objects.begin(), objects.end(), obj.name) == objects.end()) {
-          LOG(error) << "No object " << obj.name << " in map for task " << taskname;
+        if (std::find(objects.begin(), objects.end(), obj.container) == objects.end()) {
+          LOG(error) << "No container " << obj.container << " in map for task " << taskname;
           return;
         }
         auto nameHash = runtime_hash(obj.name.c_str());
@@ -334,7 +344,7 @@ AlgorithmSpec AODWriterHelpers::getOutputObjHistWriter(ConfigContext const& ctx)
         // If it's the first one, we just add it to the list.
         if (existing == inputObjects->end()) {
           obj.count = objh->mPipelineSize;
-          inputObjects->push_back(std::make_pair(key, obj));
+          inputObjects->emplace_back(key, obj);
           existing = inputObjects->end() - 1;
         } else {
           obj.count = existing->second.count;
@@ -375,43 +385,23 @@ AlgorithmSpec AODWriterHelpers::getOutputObjHistWriter(ConfigContext const& ctx)
           currentFile = filename;
         }
 
-        // translate the list-structure created by the registry into a directory structure within the file
-        std::function<void(TList*, TDirectory*)> writeListToFile;
-        writeListToFile = [&](TList* list, TDirectory* parentDir) {
-          TIter next(list);
-          TObject* object = nullptr;
-          while ((object = next())) {
-            if (object->InheritsFrom(TList::Class())) {
-              writeListToFile(static_cast<TList*>(object), parentDir->mkdir(object->GetName(), object->GetName(), true));
-            } else {
-              int objSize = parentDir->WriteObjectAny(object, object->Class(), object->GetName());
-              static int maxSizeWritten = 0;
-              if (objSize > maxSizeWritten) {
-                auto& monitoring = pc.services().get<Monitoring>();
-                maxSizeWritten = objSize;
-                monitoring.send(Metric{fmt::format("{}/{}:{}", object->ClassName(), object->GetName(), objSize), "aod-largest-object-written"}.addTag(tags::Key::Subsystem, tags::Value::DPL));
-              }
-              auto* written = list->Remove(object);
-              delete written;
-            }
+        // FIXME: handle folders
+        auto* currentDir = f[route.policy]->GetDirectory(currentDirectory.c_str());
+        // The name contains a path...
+        if (sourceType == HistogramRegistrySource) {
+          TDirectory* currentFolder = currentDir;
+          std::string objName = entry.name;
+          auto lastSlash = entry.name.rfind('/');
+          auto containerName = obj.container;
+          if (lastSlash != std::string::npos) {
+            auto dirname = entry.name.substr(0, lastSlash);
+            objName = entry.name.substr(lastSlash + 1);
+            containerName += "/" + dirname;
+            currentFolder = currentDir->mkdir(containerName.c_str(), "", kTRUE);
           }
-        };
-
-        TDirectory* currentDir = f[route.policy]->GetDirectory(currentDirectory.c_str());
-        if (route.sourceType == OutputObjSourceType::HistogramRegistrySource) {
-          auto* outputList = static_cast<TList*>(entry.obj);
-          outputList->SetOwner(false);
-
-          // if registry should live in dedicated folder a TNamed object is appended to the list
-          if (outputList->Last() && outputList->Last()->IsA() == TNamed::Class()) {
-            delete outputList->Last();
-            outputList->RemoveLast();
-            currentDir = currentDir->mkdir(outputList->GetName(), outputList->GetName(), true);
-          }
-
-          writeListToFile(outputList, currentDir);
-          outputList->SetOwner();
-          delete outputList;
+          currentFolder = currentDir->GetDirectory(containerName.c_str());
+          currentFolder->WriteObjectAny(entry.obj, entry.kind, objName.c_str());
+          delete (TObject*)entry.obj;
           entry.obj = nullptr;
         } else {
           currentDir->WriteObjectAny(entry.obj, entry.kind, entry.name.c_str());
diff --git a/Framework/Core/include/Framework/AnalysisManagers.h b/Framework/Core/include/Framework/AnalysisManagers.h
@@ -11,6 +11,7 @@
 
 #ifndef FRAMEWORK_ANALYSISMANAGERS_H
 #define FRAMEWORK_ANALYSISMANAGERS_H
+#include "DataAllocator.h"
 #include "Framework/AnalysisHelpers.h"
 #include "Framework/DataSpecUtils.h"
 #include "Framework/GroupedCombinations.h"
@@ -247,7 +248,10 @@ template <is_histogram_registry T>
 bool postRunOutput(EndOfStreamContext& context, T& hr)
 {
   auto& deviceSpec = context.services().get<o2::framework::DeviceSpec const>();
-  context.outputs().snapshot(hr.ref(deviceSpec.inputTimesliceId, deviceSpec.maxInputTimeslices), *(hr.getListOfHistograms()));
+  auto sendHistos = [deviceSpec, &context](HistogramRegistry const& self, TNamed* obj) mutable {
+    context.outputs().snapshot(self.ref(deviceSpec.inputTimesliceId, deviceSpec.maxInputTimeslices), *obj);
+  };
+  hr.apply(sendHistos);
   hr.clean();
   return true;
 }
diff --git a/Framework/Core/include/Framework/HistogramRegistry.h b/Framework/Core/include/Framework/HistogramRegistry.h
@@ -173,16 +173,15 @@ class HistogramRegistry
   template <typename T>
   std::shared_ptr<T> operator()(const HistName& histName);
 
+  // Apply @a callback on every single entry in the registry
+  void apply(std::function<void(HistogramRegistry const&, TNamed* named)> callback) const;
   // return the OutputSpec associated to the HistogramRegistry
   OutputSpec const spec();
 
-  OutputRef ref(uint16_t idx, uint16_t pipelineSize);
+  OutputRef ref(uint16_t idx, uint16_t pipelineSize) const;
 
   void setHash(uint32_t hash);
 
-  /// returns the list of histograms, properly sorted for writing.
-  TList* getListOfHistograms();
-
   /// deletes all the histograms from the registry
   void clean();
 
@@ -220,16 +219,13 @@ class HistogramRegistry
 
   // helper function to find the histogram position in the registry
   template <typename T>
-  uint32_t getHistIndex(const T& histName);
+  uint32_t getHistIndex(const T& histName) const;
 
   constexpr uint32_t imask(uint32_t i) const
   {
     return i & REGISTRY_BITMASK;
   }
 
-  // helper function to create resp. find the subList defined by path
-  TList* getSubList(TList* list, std::deque<std::string>& path);
-
   // helper function to split user defined path/to/hist/name string
   std::deque<std::string> splitPath(const std::string& pathAndNameUser);
 
@@ -431,7 +427,7 @@ std::shared_ptr<T> HistogramRegistry::operator()(const HistName& histName)
 }
 
 template <typename T>
-uint32_t HistogramRegistry::getHistIndex(const T& histName)
+uint32_t HistogramRegistry::getHistIndex(const T& histName) const
 {
   if (O2_BUILTIN_LIKELY(histName.hash == mRegistryKey[histName.idx])) {
     return histName.idx;
diff --git a/Framework/Core/include/Framework/OutputObjHeader.h b/Framework/Core/include/Framework/OutputObjHeader.h
@@ -44,6 +44,8 @@ struct OutputObjHeader : public BaseHeader {
   uint32_t mTaskHash;
   uint16_t mPipelineIndex = 0;
   uint16_t mPipelineSize = 1;
+  // Name of the actual container for the object, e.g. the HistogramRegistry name
+  char containerName[64] = {0};
 
   constexpr OutputObjHeader()
     : BaseHeader(sizeof(OutputObjHeader), sHeaderType, sSerializationMethod, sVersion),
diff --git a/Framework/Core/src/HistogramRegistry.cxx b/Framework/Core/src/HistogramRegistry.cxx
@@ -51,9 +51,12 @@ OutputSpec const HistogramRegistry::spec()
   return OutputSpec{OutputLabel{mName}, "ATSK", desc, 0, Lifetime::QA};
 }
 
-OutputRef HistogramRegistry::ref(uint16_t pipelineIndex, uint16_t pipelineSize)
+OutputRef HistogramRegistry::ref(uint16_t pipelineIndex, uint16_t pipelineSize) const
 {
-  return OutputRef{std::string{mName}, 0, o2::header::Stack{OutputObjHeader{mPolicy, OutputObjSourceType::HistogramRegistrySource, mTaskHash, pipelineIndex, pipelineSize}}};
+  OutputObjHeader header{mPolicy, OutputObjSourceType::HistogramRegistrySource, mTaskHash, pipelineIndex, pipelineSize};
+  // Copy the name of the registry to the haeder.
+  strncpy(header.containerName, mName.data(), 64);
+  return OutputRef{std::string{mName}, 0, o2::header::Stack{header}};
 }
 
 void HistogramRegistry::setHash(uint32_t hash)
@@ -282,87 +285,24 @@ void HistogramRegistry::print(bool showAxisDetails)
   LOGF(info, "");
 }
 
-// create output structure will be propagated to file-sink
-TList* HistogramRegistry::getListOfHistograms()
+void HistogramRegistry::apply(std::function<void(HistogramRegistry const&, TNamed* named)> callback) const
 {
-  TList* list = new TList();
-  list->SetName(mName.data());
-
-  if (mSortHistos) {
-    auto caseInsensitiveCompare = [](const std::string& s1, const std::string& s2) {
-      return std::lexicographical_compare(s1.begin(), s1.end(), s2.begin(), s2.end(),
-                                          [](char c1, char c2) { return std::tolower(static_cast<unsigned char>(c1)) < std::tolower(static_cast<unsigned char>(c2)); });
-    };
-    std::sort(mRegisteredNames.begin(), mRegisteredNames.end(), caseInsensitiveCompare);
-  }
-
-  for (auto& curHistName : mRegisteredNames) {
+  // Keep the list sorted as originally done to avoid hidden dependency on the order, for now , for now.
+  auto finalList = mRegisteredNames;
+  auto caseInsensitiveCompare = [](const std::string& s1, const std::string& s2) {
+    return std::lexicographical_compare(s1.begin(), s1.end(), s2.begin(), s2.end(),
+                                        [](char c1, char c2) { return std::tolower(static_cast<unsigned char>(c1)) < std::tolower(static_cast<unsigned char>(c2)); });
+  };
+  std::sort(finalList.begin(), finalList.end(), caseInsensitiveCompare);
+  for (auto& curHistName : finalList) {
     TNamed* rawPtr = nullptr;
     std::visit([&](const auto& sharedPtr) { rawPtr = (TNamed*)sharedPtr.get(); }, mRegistryValue[getHistIndex(HistName{curHistName.data()})]);
-    if (rawPtr) {
-      std::deque<std::string> path = splitPath(rawPtr->GetName());
-      std::string name = path.back();
-      path.pop_back();
-      TList* targetList{getSubList(list, path)};
-      if (targetList) {
-        rawPtr->SetName(name.data());
-        targetList->Add(rawPtr);
-      } else {
-        LOGF(fatal, "Specified subfolder could not be created.");
-      }
-    }
-  }
-
-  // place lists always at the top
-  std::function<void(TList*)> moveListsToTop;
-  moveListsToTop = [&](TList* list) {
-    TIter next(list);
-    TNamed* subList = nullptr;
-    std::vector<TObject*> subLists;
-    while ((subList = (TNamed*)next())) {
-      if (subList->InheritsFrom(TList::Class())) {
-        subLists.push_back(subList);
-        moveListsToTop((TList*)subList);
-      }
-    }
-    std::reverse(subLists.begin(), subLists.end());
-    for (auto curList : subLists) {
-      list->Remove(curList);
-      list->AddFirst(curList);
-    }
-  };
-  moveListsToTop(list);
-
-  // create dedicated directory containing all of the registrys histograms
-  if (mCreateRegistryDir) {
-    // propagate this to the writer by adding a 'flag' to the output list
-    list->AddLast(new TNamed("createFolder", ""));
-  }
-  return list;
-}
-
-// helper function to create resp. find the subList defined by path
-TList* HistogramRegistry::getSubList(TList* list, std::deque<std::string>& path)
-{
-  if (path.empty()) {
-    return list;
-  }
-  TList* targetList{nullptr};
-  std::string nextList = path[0];
-  path.pop_front();
-  if (auto subList = (TList*)list->FindObject(nextList.data())) {
-    if (subList->InheritsFrom(TList::Class())) {
-      targetList = getSubList((TList*)subList, path);
-    } else {
-      return nullptr;
+    if (!rawPtr) {
+      // Skipping empty histograms
+      continue;
     }
-  } else {
-    subList = new TList();
-    subList->SetName(nextList.data());
-    list->Add(subList);
-    targetList = getSubList(subList, path);
+    callback(*this, rawPtr);
   }
-  return targetList;
 }
 
 // helper function to split user defined path/to/hist/name string
diff --git a/Framework/TestWorkflows/src/o2TestHistograms.cxx b/Framework/TestWorkflows/src/o2TestHistograms.cxx
@@ -43,6 +43,16 @@ struct EtaAndClsHistogramsSimple {
   Configurable<std::string> trackFilterString{"track-filter", "o2::aod::track::pt < 10.f", "Track filter string"};
   Filter trackFilter = o2::aod::track::pt < 10.f;
 
+  HistogramRegistry registry{
+    "registry",
+    {
+      {"a/b/eta", "#Eta", {HistType::kTH1F, {{100, -2.0, 2.0}}}},                          //
+      {"a/phi", "#Phi", {HistType::kTH1D, {{102, 0, 2 * M_PI}}}},                          //
+      {"c/pt", "p_{T}", {HistType::kTH1D, {{1002, -0.01, 50.1}}}},                         //
+      {"ptToPt", "#ptToPt", {HistType::kTH2F, {{100, -0.01, 10.01}, {100, -0.01, 10.01}}}} //
+    } //
+  };
+
   void init(InitContext&)
   {
     if (!trackFilterString->empty()) {
@@ -56,6 +66,11 @@ struct EtaAndClsHistogramsSimple {
     for (auto& track : tracks) {
       etaClsH->Fill(track.eta(), track.pt());
       skimEx(track.pt(), track.eta());
+
+      registry.fill(HIST("a/b/eta"), track.eta());
+      registry.fill(HIST("a/phi"), track.phi());
+      registry.fill(HIST("c/pt"), track.pt());
+      registry.fill(HIST("ptToPt"), track.pt(), track.signed1Pt());
     }
   }
 };
@@ -66,6 +81,16 @@ struct EtaAndClsHistogramsIUSimple {
   Configurable<std::string> trackFilterString{"track-filter", "o2::aod::track::pt < 10.f", "Track filter string"};
   Filter trackFilter = o2::aod::track::pt < 10.f;
 
+  HistogramRegistry registry{
+    "registry",
+    {
+      {"a/b/eta", "#Eta", {HistType::kTH1F, {{100, -2.0, 2.0}}}},                          //
+      {"a/phi", "#Phi", {HistType::kTH1D, {{102, 0, 2 * M_PI}}}},                          //
+      {"c/pt", "p_{T}", {HistType::kTH1D, {{1002, -0.01, 50.1}}}},                         //
+      {"ptToPt", "#ptToPt", {HistType::kTH2F, {{100, -0.01, 10.01}, {100, -0.01, 10.01}}}} //
+    } //
+  };
+
   void init(InitContext&)
   {
     if (!trackFilterString->empty()) {
@@ -79,12 +104,28 @@ struct EtaAndClsHistogramsIUSimple {
     for (auto& track : tracks) {
       etaClsH->Fill(track.eta(), track.pt());
       skimEx(track.pt(), track.eta());
+
+      registry.fill(HIST("a/b/eta"), track.eta());
+      registry.fill(HIST("a/phi"), track.phi());
+      registry.fill(HIST("c/pt"), track.pt());
+      registry.fill(HIST("ptToPt"), track.pt(), track.signed1Pt());
     }
   }
 };
 
 struct EtaAndClsHistogramsFull {
   OutputObj<TH3F> etaClsH{TH3F("eta_vs_cls_vs_sigmapT", "#eta vs N_{cls} vs sigma_{1/pT}", 102, -2.01, 2.01, 160, -0.5, 159.5, 100, 0, 10)};
+
+  HistogramRegistry registry{
+    "registry",
+    {
+      {"a/b/eta", "#Eta", {HistType::kTH1F, {{100, -2.0, 2.0}}}},                          //
+      {"a/phi", "#Phi", {HistType::kTH1D, {{102, 0, 2 * M_PI}}}},                          //
+      {"c/pt", "p_{T}", {HistType::kTH1D, {{1002, -0.01, 50.1}}}},                         //
+      {"ptToPt", "#ptToPt", {HistType::kTH2F, {{100, -0.01, 10.01}, {100, -0.01, 10.01}}}} //
+    } //
+  };
+
   Configurable<std::string> trackFilterString{"track-filter", "o2::aod::track::pt < 10.f", "Track filter string"};
   Filter trackFilter = o2::aod::track::pt < 10.f;
 
@@ -100,6 +141,11 @@ struct EtaAndClsHistogramsFull {
     LOGP(info, "Invoking the run 3 one");
     for (auto& track : tracks) {
       etaClsH->Fill(track.eta(), track.tpcNClsFindable(), track.sigma1Pt());
+
+      registry.fill(HIST("a/b/eta"), track.eta());
+      registry.fill(HIST("a/phi"), track.phi());
+      registry.fill(HIST("c/pt"), track.pt());
+      registry.fill(HIST("ptToPt"), track.pt(), track.signed1Pt());
     }
   }
 };
