TRD Kr cluster writing to EOS

Detectors/TRD/workflow/CMakeLists.txt

@@ -21,7 +21,7 @@ o2_add_library(TRDWorkflow
                        src/EntropyDecoderSpec.cxx
                        src/EntropyEncoderSpec.cxx
                        src/TrackBasedCalibSpec.cxx
-                       src/KrClustererSpec.cxx
+                       include/TRDWorkflow/KrClustererSpec.h
                        include/TRDWorkflow/VdAndExBCalibSpec.h
                        include/TRDWorkflow/GainCalibSpec.h
                        include/TRDWorkflow/TRDPulseHeightSpec.h

Detectors/TRD/workflow/include/TRDWorkflow/KrClustererSpec.h

@@ -19,18 +19,82 @@
 // input TRD digits, TRD trigger records
 // output Kr clusters
 
+#include "TRDWorkflow/KrClustererSpec.h"
+#include "TRDCalibration/KrClusterFinder.h"
+#include "Framework/Task.h"
+#include "Framework/ConfigParamRegistry.h"
 #include "Framework/DataProcessorSpec.h"
+#include "TStopwatch.h"
+#include <fairlogger/Logger.h>
 
 using namespace o2::framework;
 
 namespace o2
 {
 namespace trd
 {
-/// create a processor spec
-framework::DataProcessorSpec getKrClustererSpec();
+class TRDKrClustererDevice : public Task
+{
+ public:
+  TRDKrClustererDevice() = default;
+  ~TRDKrClustererDevice() override = default;
+  void init(InitContext& ic) final;
+  void run(ProcessingContext& pc) final;
+  void endOfStream(framework::EndOfStreamContext& ec) final;
+
+ private:
+  o2::trd::KrClusterFinder mKrClFinder;
+};
+
+void TRDKrClustererDevice::init(InitContext& ic)
+{
+  mKrClFinder.init();
+}
+
+void TRDKrClustererDevice::run(ProcessingContext& pc)
+{
+  TStopwatch timer;
+
+  const auto digits = pc.inputs().get<gsl::span<Digit>>("digits");
+  const auto triggerRecords = pc.inputs().get<gsl::span<TriggerRecord>>("triggerRecords");
+
+  mKrClFinder.reset();
+  mKrClFinder.setInput(digits, triggerRecords);
+  timer.Start();
+  mKrClFinder.findClusters();
+  timer.Stop();
+
+  LOGP(info, "Found {} Kr clusters in {} input trigger records. Timing: CPU: {}, Real: {}",
+       mKrClFinder.getKrClusters().size(), triggerRecords.size(), timer.CpuTime(), timer.RealTime());
+
+  pc.outputs().snapshot(Output{o2::header::gDataOriginTRD, "KRCLUSTER", 0, Lifetime::Timeframe}, mKrClFinder.getKrClusters());
+  pc.outputs().snapshot(Output{o2::header::gDataOriginTRD, "TRGKRCLS", 0, Lifetime::Timeframe}, mKrClFinder.getKrTrigRecs());
+}
+
+void TRDKrClustererDevice::endOfStream(EndOfStreamContext& ec)
+{
+  LOG(info) << "Done with the cluster finding (EoS received)";
+}
+
+framework::DataProcessorSpec getKrClustererSpec()
+{
+  std::vector<InputSpec> inputs;
+  inputs.emplace_back("digits", ConcreteDataTypeMatcher{o2::header::gDataOriginTRD, "DIGITS"}, Lifetime::Timeframe);
+  inputs.emplace_back("triggerRecords", ConcreteDataTypeMatcher{o2::header::gDataOriginTRD, "TRKTRGRD"}, Lifetime::Timeframe);
+  std::vector<OutputSpec> outputs;
+  outputs.emplace_back(o2::header::gDataOriginTRD, "KRCLUSTER", 0, Lifetime::Timeframe);
+  outputs.emplace_back(o2::header::gDataOriginTRD, "TRGKRCLS", 0, Lifetime::Timeframe);
+
+  return DataProcessorSpec{
+    "trd-kr-clusterer",
+    inputs,
+    outputs,
+    AlgorithmSpec{adaptFromTask<o2::trd::TRDKrClustererDevice>()},
+    Options{}};
+}
 
 } // namespace trd
+
 } // namespace o2
 
 #endif // O2_TRD_KRCLUSTERERSPEC_H

Detectors/TRD/workflow/io/CMakeLists.txt

@@ -22,7 +22,7 @@ o2_add_library(TRDWorkflowIO
                        src/TRDTrackReaderSpec.cxx
                        src/TRDCalibWriterSpec.cxx
                        src/TRDPHReaderSpec.cxx
-                       src/KrClusterWriterSpec.cxx
+                       include/TRDWorkflowIO/KrClusterWriterSpec.h
                PUBLIC_LINK_LIBRARIES O2::DataFormatsTRD O2::SimulationDataFormat O2::DPLUtils O2::GPUDataTypeHeaders O2::DataFormatsTPC)
 
 

Detectors/TRD/workflow/io/include/TRDWorkflowIO/KrClusterWriterSpec.h

@@ -12,20 +12,170 @@
 #ifndef O2_TRD_KRWRITERSPEC_H
 #define O2_TRD_KRWRITERSPEC_H
 
+#include "Framework/DataProcessorSpec.h"
+#include "Framework/DataTakingContext.h"
+#include "DataFormatsTRD/KrCluster.h"
+#include "DataFormatsTRD/KrClusterTriggerRecord.h"
+#include "CommonUtils/MemFileHelper.h"
+#include "DetectorsCommonDataFormats/FileMetaData.h"
+
+#include <TFile.h>
+#include <TTree.h>
+#include <filesystem>
+#include <fmt/format.h>
+
 namespace o2
 {
-namespace framework
+namespace trd
 {
-struct DataProcessorSpec;
-}
-} // namespace o2
 
-namespace o2
+class TRDKrClsWriterTask : public o2::framework::Task
 {
-namespace trd
+ public:
+  TRDKrClsWriterTask() = default;
+
+  void init(o2::framework::InitContext& ic) final
+  {
+    mOutputDir = o2::utils::Str::rectifyDirectory(ic.options().get<std::string>("output-dir"));
+
+    // should we write meta files for epn2eos?
+    mMetaFileDir = ic.options().get<std::string>("meta-output-dir");
+    if (mMetaFileDir != "/dev/null") {
+      mMetaFileDir = o2::utils::Str::rectifyDirectory(mMetaFileDir);
+      mStoreMetaFile = true;
+    }
+
+    LOGP(info, "Storing output in {}, meta file writing enabled: {}", mOutputDir, mStoreMetaFile);
+    mAutoSave = ic.options().get<int>("autosave-interval");
+
+    char hostname[_POSIX_HOST_NAME_MAX];
+    gethostname(hostname, _POSIX_HOST_NAME_MAX);
+    mHostName = hostname;
+    mHostName = mHostName.substr(0, mHostName.find('.'));
+  }
+
+  void createOutputFile(int runNumber, o2::framework::ProcessingContext& pc)
+  {
+    mFileName = fmt::format("o2_trdKrCls_run{}_{}.root", runNumber, mHostName);
+    auto fileNameTmp = o2::utils::Str::concat_string(mOutputDir, mFileName, ".part");
+    mFileOut = std::make_unique<TFile>(fileNameTmp.c_str(), "recreate");
+    mTreeOut = std::make_unique<TTree>("krData", "TRD krypton cluster data");
+    mTreeOut->Branch("cluster", &krClusterPtr);
+    mTreeOut->Branch("trigRec", &krTrigRecPtr);
+    mDataTakingContext = pc.services().get<DataTakingContext>();
+    mOutputFileCreated = true;
+  }
+
+  void writeToFile()
+  {
+    if (!mOutputFileCreated) {
+      return;
+    }
+    mFileOut->cd();
+    mTreeOut->Write();
+  }
+
+  void closeOutputFile()
+  {
+    if (!mOutputFileCreated) {
+      return;
+    }
+    writeToFile();
+    mTreeOut.reset();
+    mFileOut->Close();
+    mFileOut.reset();
+    auto fileNameWithPath = mOutputDir + mFileName;
+    std::filesystem::rename(o2::utils::Str::concat_string(mOutputDir, mFileName, ".part"), fileNameWithPath);
+    if (mStoreMetaFile) {
+      o2::dataformats::FileMetaData fileMetaData; // object with information for meta data file
+      fileMetaData.fillFileData(fileNameWithPath);
+      fileMetaData.setDataTakingContext(mDataTakingContext);
+      fileMetaData.type = "calib";
+      fileMetaData.priority = "high";
+      auto metaFileNameTmp = fmt::format("{}{}.tmp", mMetaFileDir, mFileName);
+      auto metaFileName = fmt::format("{}{}.done", mMetaFileDir, mFileName);
+      try {
+        std::ofstream metaFileOut(metaFileNameTmp);
+        metaFileOut << fileMetaData;
+        metaFileOut.close();
+        std::filesystem::rename(metaFileNameTmp, metaFileName);
+      } catch (std::exception const& e) {
+        LOG(error) << "Failed to store meta data file " << metaFileName << ", reason: " << e.what();
+      }
+    }
+  }
+
+  void run(o2::framework::ProcessingContext& pc) final
+  {
+    if (!mOutputFileCreated) {
+      auto tInfo = pc.services().get<o2::framework::TimingInfo>();
+      createOutputFile(tInfo.runNumber, pc);
+    }
+    if (pc.transitionState() == TransitionHandlingState::Requested) {
+      LOG(info) << "Run stop requested, closing output file";
+      mRunStopRequested = true;
+      closeOutputFile();
+    }
+    if (mRunStopRequested) {
+      return;
+    }
+    auto cluster = pc.inputs().get<gsl::span<KrCluster>>("krcluster");
+    auto triggerRecords = pc.inputs().get<gsl::span<KrClusterTriggerRecord>>("krtrigrec");
+    for (const auto& cls : cluster) {
+      krCluster.push_back(cls);
+    }
+    for (const auto& trig : triggerRecords) {
+      krTrigRec.push_back(trig);
+    }
+    mTreeOut->Fill();
+    krCluster.clear();
+    krTrigRec.clear();
+    if (mAutoSave > 0 && ++mTFCounter % mAutoSave == 0) {
+      writeToFile();
+    }
+  }
+
+  void endOfStream(o2::framework::EndOfStreamContext& ec) final
+  {
+    if (mRunStopRequested) {
+      return;
+    }
+    LOG(info) << "End of stream received, closing output file";
+    closeOutputFile();
+  }
+
+ private:
+  bool mRunStopRequested{false};
+  bool mStoreMetaFile{false};
+  bool mOutputFileCreated{false};
+  int mAutoSave{0};
+  uint64_t mTFCounter{0};
+  std::string mOutputDir{"none"};
+  std::string mMetaFileDir{"/dev/null"};
+  std::string mHostName{};
+  std::string mFileName{};
+  std::unique_ptr<TFile> mFileOut{};
+  std::unique_ptr<TTree> mTreeOut{};
+  std::vector<KrCluster> krCluster, *krClusterPtr{&krCluster};
+  std::vector<KrClusterTriggerRecord> krTrigRec, *krTrigRecPtr{&krTrigRec};
+  o2::framework::DataTakingContext mDataTakingContext;
+};
+
+framework::DataProcessorSpec getKrClusterWriterSpec()
 {
+  std::vector<InputSpec> inputs;
+  inputs.emplace_back("krcluster", "TRD", "KRCLUSTER");
+  inputs.emplace_back("krtrigrec", "TRD", "TRGKRCLS");
 
-o2::framework::DataProcessorSpec getKrClusterWriterSpec();
+  return DataProcessorSpec{"kr-cluster-writer",
+                           inputs,
+                           Outputs{},
+                           AlgorithmSpec{adaptFromTask<o2::trd::TRDKrClsWriterTask>()},
+                           Options{
+                             {"output-dir", VariantType::String, "none", {"Output directory for data. Defaults to current working directory"}},
+                             {"meta-output-dir", VariantType::String, "/dev/null", {"metadata output directory, must exist (if not /dev/null)"}},
+                             {"autosave-interval", VariantType::Int, 0, {"Write output to file for every n-th TF. 0 means this feature is OFF"}}}};
+}
 
 } // end namespace trd
 } // end namespace o2