First version of CCDB implementation

Author: ChSonnabend

GPU/GPUTracking/Definitions/GPUSettingsList.h

@@ -249,6 +249,14 @@ AddOption(nnClassificationPath, std::string, "network_class.onnx", "", 0, "The c
 AddOption(nnClassThreshold, float, 0.5, "", 0, "The cutoff at which clusters will be accepted / rejected.")
 AddOption(nnRegressionPath, std::string, "network_reg.onnx", "", 0, "The regression network path")
 AddOption(nnSigmoidTrafoClassThreshold, int, 1, "", 0, "If true (default), then the classification threshold is transformed by an inverse sigmoid function. This depends on how the network was trained (with a sigmoid as acitvation function in the last layer or not).")
+// CCDB
+AddOption(nnLoadFromCCDB, int, 1, "", 0, "If 1 networks are fetched from ccdb, else locally")
+AddOption(nnCCDBURL, std::string, "http://alice-ccdb.cern.ch", "", 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(nnCCDBLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN")
+AddOption(nnCCDBBeamType, std::string, "", "", 0, "Distinguishes between networks trained for different beam types. Options: PbPb, pp")
+AddOption(nnCCDBInteractionRate, int, -1, "", 0, "Distinguishes between networks for different interaction rates [kHz].")
 AddHelp("help", 'h')
 EndConfig()
 

GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx

@@ -615,6 +615,27 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
   const GPUSettingsProcessingNNclusterizer& nn_settings = GetProcessingSettings().nn;
   GPUTPCNNClusterizerHost nnApplication; // potentially this needs to be GPUTPCNNClusterizerHost nnApplication[NSECTORS]; Technically ONNX ->Run() is threadsafe at inference time since its read-only
   if (GetProcessingSettings().nn.applyNNclusterizer) {
+    if(nn_settings.nnLoadFromCCDB) {
+      std::unordered_map<std::string, std::string> ccdbSettings = {
+        {"nnCCDBPath", nn_settings.nnCCDBPath},
+        {"inputDType", nn_settings.inputDType},
+        {"outputDType", nn_settings.outputDType},
+        {"nnCCDBWithMomentum", std::to_string(nn_settings.nnCCDBWithMomentum)},
+        {"nnCCDBLayerType", nn_settings.nnCCDBLayerType},
+        {"nnCCDBBeamType", nn_settings.nnCCDBBeamType},
+        {"nnCCDBInteractionRate", std::to_string(nn_settings.nnCCDBInteractionRate)}
+      };
+
+      std::unordered_map<std::string, std::string> networkRetrieval = ccdbSettings;
+
+      networkRetrieval["nnCCDBEvalType"] = "classification_c1";
+      networkRetrieval["outputFile"] = "net_classification_c1.onnx";
+      nnApplication.loadFromCCDB(networkRetrieval);
+
+      networkRetrieval["nnCCDBEvalType"] = "regression_c1";
+      networkRetrieval["outputFile"] = "net_regression_c1.onnx";
+      nnApplication.loadFromCCDB(networkRetrieval);
+    }
     uint32_t maxClusters = 0;
     nnApplication.init(nn_settings);
     for (uint32_t iSector = 0; iSector < NSECTORS; iSector++) {

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx

@@ -26,6 +26,32 @@ GPUTPCNNClusterizerHost::GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNcl
   init(settings);
 }
 
+void GPUTPCNNClusterizerHost::loadFromCCDB(std::unordered_map<std::string, std::string> settings) {
+  o2::ccdb::CcdbApi ccdbApi;
+  ccdbApi.init(settings["nnCCDBURL"]);
+
+  metadata[settings["inputDType"]] = settings["inputDType"];
+  metadata[settings["outputDType"]] = settings["outputDType"];
+  metadata[settings["nnCCDBEvalType"]] = settings["nnCCDBEvalType"]; // classification_1C, classification_2C, regression_1C, regression_2C
+  metadata[settings["nnCCDBWithMomentum"]] = std::stoi(settings["nnCCDBWithMomentum"]); // 0, 1 -> Only for regression model
+  metadata[settings["nnCCDBLayerType"]] = settings["nnCCDBLayerType"]; // FC, CNN
+  if (settings["nnCCDBInteractionRate"] != "" && std::stoi(settings["nnCCDBInteractionRate"]) > 0) {
+    metadata[settings["nnCCDBInteractionRate"]] = settings["nnCCDBInteractionRate"];
+  }
+  if (settings["nnCCDBBeamType"] != "") {
+    metadata[settings["nnCCDBBeamType"]] = settings["nnCCDBBeamType"];
+  }
+
+  bool retrieveSuccess = ccdbApi.retrieveBlob(settings["nnPathCCDB"], ".", metadata, 1, false, settings["outputFile"]);
+  // headers = ccdbApi.retrieveHeaders(nnPathCCDB, metadata, ccdbTimestamp); // potentially needed to init some local variables
+
+  if (retrieveSuccess) {
+    LOG(info) << "Network " << settings["nnPathCCDB"] << " retrieved from CCDB, stored at " << settings["networkPathLocal"];
+  } else {
+    LOG(error) << "Failed to retrieve network from CCDB";
+  }
+}
+
 void GPUTPCNNClusterizerHost::init(const GPUSettingsProcessingNNclusterizer& settings)
 {
   OrtOptions = {

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.h

@@ -41,6 +41,7 @@ class GPUTPCNNClusterizerHost
 
   void init(const GPUSettingsProcessingNNclusterizer&);
   void initClusterizer(const GPUSettingsProcessingNNclusterizer&, GPUTPCNNClusterizer&);
+  void loadFromCCDB(std::unordered_map<std::string, std::string>);
 
   void networkInference(o2::ml::OrtModel model, GPUTPCNNClusterizer& clusterer, size_t size, float* output, int32_t dtype);