Working version of NN CCDB fetching and loading to file

ChSonnabend · ChSonnabend · commit 3775044da0ba · 2025-11-21T13:02:54.000+01:00
diff --git a/Common/ML/include/ML/OrtInterface.h b/Common/ML/include/ML/OrtInterface.h
@@ -51,6 +51,7 @@ class OrtModel
   void initOptions(std::unordered_map<std::string, std::string> optionsMap);
   void initEnvironment();
   void initSession();
+  void initSessionFromBuffer(const void* buffer, size_t bufferSize);
   void memoryOnDevice(int32_t = 0);
   bool isInitialized() { return mInitialized; }
   void resetSession();
diff --git a/Common/ML/src/OrtInterface.cxx b/Common/ML/src/OrtInterface.cxx
@@ -138,6 +138,21 @@ void OrtModel::initEnvironment()
   (mPImplOrt->env)->DisableTelemetryEvents(); // Disable telemetry events
 }
 
+void OrtModel::initSessionFromBuffer(const void* buffer, size_t bufferSize)
+{
+  mPImplOrt->session = std::make_unique<Ort::Session>(*mPImplOrt->env,
+                                                      static_cast<const uint8_t*>(buffer),
+                                                      bufferSize,
+                                                      mPImplOrt->sessionOptions);
+  mPImplOrt->ioBinding = std::make_unique<Ort::IoBinding>(*mPImplOrt->session);
+
+  setIO();
+
+  if (mLoggingLevel < 2) {
+    LOG(info) << "(ORT) Model loaded successfully from buffer! (inputs: " << printShape(mInputShapes, mInputNames) << ", outputs: " << printShape(mOutputShapes, mInputNames) << ")";
+  }
+}
+
 void OrtModel::initSession()
 {
   if (mAllocateDeviceMemory) {
diff --git a/Detectors/TPC/calibration/src/NeuralNetworkClusterizer.cxx b/Detectors/TPC/calibration/src/NeuralNetworkClusterizer.cxx
@@ -37,6 +37,8 @@ void NeuralNetworkClusterizer::loadIndividualFromCCDB(std::map<std::string, std:
     metadata["nnCCDBBeamType"] = settings["nnCCDBBeamType"];
   }
 
+  LOG(info) << "(NN CLUS) Retrieving network " << settings["nnCCDBPath"] << " from CCDB (NeuralNetworkClusterizer.cxx)";
+
   bool retrieveSuccess = ccdbApi.retrieveBlob(settings["nnCCDBPath"], settings["outputFolder"], metadata, 1, false, settings["outputFile"]);
   // headers = ccdbApi.retrieveHeaders(settings["nnPathCCDB"], metadata, 1); // potentially needed to init some local variables
 
diff --git a/GPU/GPUTracking/Definitions/GPUSettingsList.h b/GPU/GPUTracking/Definitions/GPUSettingsList.h
@@ -285,14 +285,14 @@ AddOption(nnClusterizerUseClassification, int, 1, "", 0, "If 1, the classificati
 AddOption(nnClusterizerForceGpuInputFill, int, 0, "", 0, "Forces to use the fillInputNNGPU function")
 // CCDB
 AddOption(nnLoadFromCCDB, int, 0, "", 0, "If 1 networks are fetched from ccdb, else locally")
+AddOption(nnCCDBDumpToFile, int, 1, "", 0, "If 1, additionally dump fetched CCDB networks to nnLocalFolder")
 AddOption(nnLocalFolder, std::string, ".", "", 0, "Local folder in which the networks will be fetched")
-AddOption(nnCCDBURL, std::string, "http://ccdb-test.cern.ch:8080", "", 0, "The CCDB URL from where the network files are fetched")
 AddOption(nnCCDBPath, std::string, "Users/c/csonnabe/TPC/Clusterization", "", 0, "Folder path containing the networks")
-AddOption(nnCCDBWithMomentum, int, 1, "", 0, "Distinguishes between the network with and without momentum output for the regression")
+AddOption(nnCCDBWithMomentum, std::string, "", "", 0, "Distinguishes between the network with and without momentum output for the regression")
 AddOption(nnCCDBClassificationLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")
-AddOption(nnCCDBRegressionLayerType, std::string, "CNN", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")
-AddOption(nnCCDBBeamType, std::string, "PbPb", "", 0, "Distinguishes between networks trained for different beam types. Options: PbPb, pp")
-AddOption(nnCCDBInteractionRate, int, 50, "", 0, "Distinguishes between networks for different interaction rates [kHz].")
+AddOption(nnCCDBRegressionLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")
+AddOption(nnCCDBBeamType, std::string, "PbPb", "", 0, "Distinguishes between networks trained for different beam types. Options: PbPb, pp, pPb")
+AddOption(nnCCDBInteractionRate, std::string, "500", "", 0, "Distinguishes between networks for different interaction rates [kHz].")
 AddHelp("help", 'h')
 EndConfig()
 
diff --git a/GPU/Workflow/include/GPUWorkflow/GPUWorkflowSpec.h b/GPU/Workflow/include/GPUWorkflow/GPUWorkflowSpec.h
@@ -135,6 +135,9 @@ class GPURecoWorkflowSpec : public o2::framework::Task
     bool tpcTriggerHandling = false;
     bool isITS3 = false;
     bool useFilteredOutputSpecs = false;
+
+    // NN clusterizer
+    bool nnLoadFromCCDB = false;
   };
 
   GPURecoWorkflowSpec(CompletionPolicyData* policyData, Config const& specconfig, std::vector<int32_t> const& tpcsectors, uint64_t tpcSectorMask, std::shared_ptr<o2::base::GRPGeomRequest>& ggr, std::function<bool(o2::framework::DataProcessingHeader::StartTime)>** gPolicyOrder = nullptr);
@@ -230,7 +233,7 @@ class GPURecoWorkflowSpec : public o2::framework::Task
   uint32_t mNextThreadIndex = 0;
   bool mUpdateGainMapCCDB = true;
   std::unique_ptr<o2::gpu::GPUSettingsTF> mTFSettings;
-  std::unique_ptr<o2::gpu::GPUSettingsProcessingNNclusterizer> mNNClusterizerSettings;
+  std::map<std::string, std::string> nnCCDBSettings;
 
   Config mSpecConfig;
   std::shared_ptr<o2::base::GRPGeomRequest> mGGR;
diff --git a/GPU/Workflow/src/GPUWorkflowSpec.cxx b/GPU/Workflow/src/GPUWorkflowSpec.cxx
@@ -133,50 +133,6 @@ void GPURecoWorkflowSpec::init(InitContext& ic)
 {
   GRPGeomHelper::instance().setRequest(mGGR);
   GPUO2InterfaceConfiguration& config = *mConfig.get();
-  GPUSettingsProcessingNNclusterizer& mNNClusterizerSettings = mConfig->configProcessing.nn;
-
-  if (mNNClusterizerSettings.nnLoadFromCCDB) {
-    LOG(info) << "Loading neural networks from CCDB";
-    o2::tpc::NeuralNetworkClusterizer nnClusterizerFetcher;
-    nnClusterizerFetcher.initCcdbApi(mNNClusterizerSettings.nnCCDBURL);
-    std::map<std::string, std::string> ccdbSettings = {
-      {"nnCCDBURL", mNNClusterizerSettings.nnCCDBURL},
-      {"nnCCDBPath", mNNClusterizerSettings.nnCCDBPath},
-      {"inputDType", mNNClusterizerSettings.nnInferenceInputDType},
-      {"outputDType", mNNClusterizerSettings.nnInferenceOutputDType},
-      {"outputFolder", mNNClusterizerSettings.nnLocalFolder},
-      {"nnCCDBPath", mNNClusterizerSettings.nnCCDBPath},
-      {"nnCCDBWithMomentum", std::to_string(mNNClusterizerSettings.nnCCDBWithMomentum)},
-      {"nnCCDBBeamType", mNNClusterizerSettings.nnCCDBBeamType},
-      {"nnCCDBInteractionRate", std::to_string(mNNClusterizerSettings.nnCCDBInteractionRate)}};
-
-    std::string nnFetchFolder = mNNClusterizerSettings.nnLocalFolder;
-    std::vector<std::string> evalMode = o2::utils::Str::tokenize(mNNClusterizerSettings.nnEvalMode, ':');
-
-    if (evalMode[0] == "c1") {
-      ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings.nnCCDBClassificationLayerType;
-      ccdbSettings["nnCCDBEvalType"] = "classification_c1";
-      ccdbSettings["outputFile"] = "net_classification_c1.onnx";
-      nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings);
-    } else if (evalMode[0] == "c2") {
-      ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings.nnCCDBClassificationLayerType;
-      ccdbSettings["nnCCDBEvalType"] = "classification_c2";
-      ccdbSettings["outputFile"] = "net_classification_c2.onnx";
-      nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings);
-    }
-
-    ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings.nnCCDBRegressionLayerType;
-    ccdbSettings["nnCCDBEvalType"] = "regression_c1";
-    ccdbSettings["outputFile"] = "net_regression_c1.onnx";
-    nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings);
-    if (evalMode[1] == "r2") {
-      ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings.nnCCDBRegressionLayerType;
-      ccdbSettings["nnCCDBEvalType"] = "regression_c2";
-      ccdbSettings["outputFile"] = "net_regression_c2.onnx";
-      nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings);
-    }
-    LOG(info) << "Neural network loading done!";
-  }
 
   // Create configuration object and fill settings
   mConfig->configGRP.solenoidBzNominalGPU = 0;
@@ -185,6 +141,7 @@ void GPURecoWorkflowSpec::init(InitContext& ic)
   mTFSettings->simStartOrbit = hbfu.getFirstIRofTF(o2::InteractionRecord(0, hbfu.orbitFirstSampled)).orbit;
 
   *mConfParam = mConfig->ReadConfigurableParam();
+
   if (mConfParam->display) {
     mDisplayFrontend.reset(GPUDisplayFrontendInterface::getFrontend(mConfig->configDisplay.displayFrontend.c_str()));
     mConfig->configProcessing.eventDisplay = mDisplayFrontend.get();
@@ -814,6 +771,68 @@ void GPURecoWorkflowSpec::run(ProcessingContext& pc)
 
   // ------------------------------ Actual processing ------------------------------
 
+  if (mSpecConfig.nnLoadFromCCDB) {
+    LOG(info) << "(NN CLUS) Fetching CCDB calib objects";
+
+    auto dumpOnnxToFile = [](const char* buffer, std::size_t size, const std::string& path) {
+      const char* marker = "Accept-Ranges";
+      const char* pos = std::search(buffer, buffer + size, marker, marker + std::strlen(marker));
+
+      // Compute the actual number of bytes to write
+      std::size_t writeSize = (pos != buffer + size)
+                                  ? static_cast<std::size_t>(pos - buffer)
+                                  : size;
+
+      std::ofstream out(path, std::ios::binary | std::ios::trunc);
+      if (!out.is_open()) {
+          throw std::runtime_error("Failed to open ONNX output file: " + path);
+      }
+
+      out.write(buffer, static_cast<std::streamsize>(writeSize));
+      if (!out) {
+          throw std::runtime_error("Failed while writing ONNX data to: " + path);
+      }
+    };
+
+    GPUSettingsProcessingNNclusterizer& nnClusterizerSettings = mConfig->configProcessing.nn;
+    std::vector<std::string> evalMode = o2::utils::Str::tokenize(nnClusterizerSettings.nnEvalMode, ':');
+
+    DataRef m;
+    if (evalMode[0] == "c1") {
+      m = pc.inputs().get("nn_classification_c1");
+      const char* buffer = const_cast<char*>(m.payload);
+      size_t size = DataRefUtils::getPayloadSize(m);
+      if (nnClusterizerSettings.nnCCDBDumpToFile == 1) {
+        dumpOnnxToFile(buffer, size, "net_classification_c1.onnx");
+        LOG(info) << "(NN CLUS) Dumped nn_classification_c1 from CCDB to net_classification_c1.onnx";
+      }
+    } else if (evalMode[0] == "c2") {
+      m = pc.inputs().get("nn_classification_c2");
+      const char* buffer = const_cast<char*>(m.payload);
+      size_t size = DataRefUtils::getPayloadSize(m);
+      if (nnClusterizerSettings.nnCCDBDumpToFile == 1) {
+        dumpOnnxToFile(buffer, size, "net_classification_c2.onnx");
+        LOG(info) << "(NN CLUS) Dumped nn_classification_c2 from CCDB to net_classification_c2.onnx";
+      }
+    }
+
+    m = pc.inputs().get("nn_regression_c1");
+    const char* buffer = const_cast<char*>(m.payload);
+    size_t size = DataRefUtils::getPayloadSize(m);
+    if (nnClusterizerSettings.nnCCDBDumpToFile == 1) {
+      dumpOnnxToFile(buffer, size, "net_regression_c1.onnx");
+      LOG(info) << "(NN CLUS) Dumped nn_regression_c1 from CCDB to net_regression_c1.onnx";
+    }
+    if (evalMode[1] == "r2") {
+      m = pc.inputs().get("nn_regression_c2");
+      const char* buffer = const_cast<char*>(m.payload);
+      size_t size = DataRefUtils::getPayloadSize(m);
+      if (nnClusterizerSettings.nnCCDBDumpToFile == 1) {
+        dumpOnnxToFile(buffer, size, "net_regression_c2.onnx");
+        LOG(info) << "(NN CLUS) Dumped nn_regression_c2 from CCDB to net_regression_c2.onnx";
+      }
+    }
+  }
   if ((int32_t)(ptrs.tpcZS != nullptr) + (int32_t)(ptrs.tpcPackedDigits != nullptr && (ptrs.tpcZS == nullptr || ptrs.tpcPackedDigits->tpcDigitsMC == nullptr)) + (int32_t)(ptrs.clustersNative != nullptr) + (int32_t)(ptrs.tpcCompressedClusters != nullptr) != 1) {
     throw std::runtime_error("Invalid input for gpu tracking");
   }
@@ -1262,6 +1281,67 @@ Inputs GPURecoWorkflowSpec::inputs()
     }
   }
 
+  // NN clusterizer
+  *mConfParam = mConfig->ReadConfigurableParam();
+  if (mConfig->configProcessing.nn.nnLoadFromCCDB) {
+
+    LOG(info) << "(NN CLUS) Enabling fetching of TPC NN clusterizer from CCDB";
+    mSpecConfig.nnLoadFromCCDB = true;
+    GPUSettingsProcessingNNclusterizer& nnClusterizerSettings = mConfig->configProcessing.nn;
+
+    std::map<std::string, std::string> metadata;
+    metadata["inputDType"] = nnClusterizerSettings.nnInferenceInputDType;                               // FP16 or FP32
+    metadata["outputDType"] = nnClusterizerSettings.nnInferenceOutputDType;                             // FP16 or FP32
+    metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum;          // 0, 1 -> Only for regression model
+    metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBClassificationLayerType;                  // FC, CNN
+    metadata["nnCCDBInteractionRate"] = nnClusterizerSettings.nnCCDBInteractionRate;    // in kHz
+    metadata["nnCCDBBeamType"] = nnClusterizerSettings.nnCCDBBeamType;                                  // pp, pPb, PbPb
+
+    auto convert_map_to_metadata = [](const std::map<std::string, std::string>& inputMap, std::vector<o2::framework::CCDBMetadata>& outputMetadata) {
+      for (const auto& [key, value] : inputMap) {
+        if (value != "") {
+          outputMetadata.push_back({key, value});
+        }
+      }
+    };
+
+    std::vector<std::string> evalMode = o2::utils::Str::tokenize(nnClusterizerSettings.nnEvalMode, ':');
+    std::vector<o2::framework::CCDBMetadata> ccdb_metadata;
+
+    auto printSettings = [](const std::map<std::string, std::string>& settings) {
+      LOG(info) << "(NN CLUS) NN Clusterizer CCDB settings:";
+      for (const auto& [key, value] : settings) {
+        LOG(info) << "  " << key << " : " << value;
+      }
+    };
+    printSettings(metadata);
+
+    if (evalMode[0] == "c1") {
+      metadata["nnCCDBEvalType"] = "classification_c1";
+      convert_map_to_metadata(metadata, ccdb_metadata);
+      inputs.emplace_back("nn_classification_c1", "TPC", "NNCLUSTERIZER_C1", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath, ccdb_metadata, 0));
+      LOG(info) << "(NN CLUS) Loading NN clusterizer classification (c1) from CCDB";
+    } else if (evalMode[0] == "c2") {
+      metadata["nnCCDBEvalType"] = "classification_c2";
+      convert_map_to_metadata(metadata, ccdb_metadata);
+      inputs.emplace_back("nn_classification_c2", "TPC", "NNCLUSTERIZER_C2", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath, ccdb_metadata, 0));
+      LOG(info) << "(NN CLUS) Loading NN clusterizer classification (c2) from CCDB";
+    }
+
+    metadata["nnCCDBEvalType"] = "regression_c1";
+    metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBRegressionLayerType;
+    convert_map_to_metadata(metadata, ccdb_metadata);
+    inputs.emplace_back("nn_regression_c1", "TPC", "NNCLUSTERIZER_R1", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath, ccdb_metadata, 0));
+    LOG(info) << "(NN CLUS) Loading NN clusterizer regression (r1) from CCDB";
+
+    if (evalMode[1] == "r2") {
+      metadata["nnCCDBEvalType"] = "regression_c2";
+      convert_map_to_metadata(metadata, ccdb_metadata);
+      inputs.emplace_back("nn_regression_c2", "TPC", "NNCLUSTERIZER_R2", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath, ccdb_metadata, 0));
+      LOG(info) << "(NN CLUS) Loading NN clusterizer regression (r2) from CCDB";
+    }
+  }
+
   return inputs;
 };
 

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,8 @@ void NeuralNetworkClusterizer::loadIndividualFromCCDB(std::map<std::string, std:`
`37`	`37`	`metadata["nnCCDBBeamType"] = settings["nnCCDBBeamType"];`
`38`	`38`	`}`
`39`	`39`
	`40`	`+ LOG(info) << "(NN CLUS) Retrieving network " << settings["nnCCDBPath"] << " from CCDB (NeuralNetworkClusterizer.cxx)";`
	`41`	`+`
`40`	`42`	`bool retrieveSuccess = ccdbApi.retrieveBlob(settings["nnCCDBPath"], settings["outputFolder"], metadata, 1, false, settings["outputFile"]);`
`41`	`43`	`// headers = ccdbApi.retrieveHeaders(settings["nnPathCCDB"], metadata, 1); // potentially needed to init some local variables`
`42`	`44`