Working CCDB API calls (tested with test-ccdb)

Author: ChSonnabend

GPU/GPUTracking/Definitions/GPUSettingsList.h

@@ -229,6 +229,8 @@ AddOption(nnInferenceDevice, std::string, "CPU", "", 0, "(std::string) Specify i
 AddOption(nnInferenceDeviceId, unsigned int, 0, "", 0, "(unsigned int) Specify inference device id")
 AddOption(nnInferenceAllocateDevMem, int, 0, "", 0, "(bool, default = 0), if the device memory should be allocated for inference")
 AddOption(nnInferenceDtype, std::string, "fp32", "", 0, "(std::string) Specify the datatype for which inference is performed (fp32: default, fp16)") // fp32 or fp16
+AddOption(nnInferenceInputDType, std::string, "FP32", "", 0, "(std::string) Specify the datatype for which inference is performed (FP32: default, fp16)") // fp32 or fp16
+AddOption(nnInferenceOutputDType, std::string, "FP32", "", 0, "(std::string) Specify the datatype for which inference is performed (fp32: default, fp16)") // fp32 or fp16
 AddOption(nnInferenceIntraOpNumThreads, int, 1, "", 0, "Number of threads used to evaluate one neural network (ONNX: SetIntraOpNumThreads). 0 = auto-detect, can lead to problems on SLURM systems.")
 AddOption(nnInferenceInterOpNumThreads, int, 1, "", 0, "Number of threads used to evaluate one neural network (ONNX: SetInterOpNumThreads). 0 = auto-detect, can lead to problems on SLURM systems.")
 AddOption(nnInferenceEnableOrtOptimization, unsigned int, 99, "", 0, "Enables graph optimizations in ONNX Runtime. Can be [0, 1, 2, 99] -> see https://github.com/microsoft/onnxruntime/blob/3f71d637a83dc3540753a8bb06740f67e926dc13/include/onnxruntime/core/session/onnxruntime_c_api.h#L347")
@@ -251,12 +253,13 @@ AddOption(nnRegressionPath, std::string, "network_reg.onnx", "", 0, "The regress
 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(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(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].")
+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].")
 AddHelp("help", 'h')
 EndConfig()
 

GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx

@@ -616,26 +616,29 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
   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 = {
+      std::map<std::string, std::string> ccdbSettings = {
+        {"nnCCDBURL", nn_settings.nnCCDBURL},
         {"nnCCDBPath", nn_settings.nnCCDBPath},
-        {"inputDType", nn_settings.inputDType},
-        {"outputDType", nn_settings.outputDType},
+        {"inputDType", nn_settings.nnInferenceInputDType},
+        {"outputDType", nn_settings.nnInferenceOutputDType},
         {"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;
+      std::map<std::string, std::string> networkRetrieval = ccdbSettings;
 
+      networkRetrieval["nnCCDBLayerType"] = nn_settings.nnCCDBClassificationLayerType;
       networkRetrieval["nnCCDBEvalType"] = "classification_c1";
       networkRetrieval["outputFile"] = "net_classification_c1.onnx";
       nnApplication.loadFromCCDB(networkRetrieval);
 
+      networkRetrieval["nnCCDBLayerType"] = nn_settings.nnCCDBRegressionLayerType;
       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

@@ -16,6 +16,7 @@
 
 #include "GPUTPCNNClusterizerHost.h"
 #include "GPUTPCNNClusterizer.h"
+#include "CCDB/CcdbApi.h"
 #include "GPUSettings.h"
 #include "ML/3rdparty/GPUORTFloat16.h"
 
@@ -26,27 +27,27 @@ GPUTPCNNClusterizerHost::GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNcl
   init(settings);
 }
 
-void GPUTPCNNClusterizerHost::loadFromCCDB(std::unordered_map<std::string, std::string> settings) {
+void GPUTPCNNClusterizerHost::loadFromCCDB(std::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
+  metadata["inputDType"] = settings["inputDType"];
+  metadata["outputDType"] = settings["outputDType"];
+  metadata["nnCCDBEvalType"] = settings["nnCCDBEvalType"]; // classification_1C, classification_2C, regression_1C, regression_2C
+  metadata["nnCCDBWithMomentum"] = settings["nnCCDBWithMomentum"]; // 0, 1 -> Only for regression model
+  metadata["nnCCDBLayerType"] = settings["nnCCDBLayerType"]; // FC, CNN
   if (settings["nnCCDBInteractionRate"] != "" && std::stoi(settings["nnCCDBInteractionRate"]) > 0) {
-    metadata[settings["nnCCDBInteractionRate"]] = settings["nnCCDBInteractionRate"];
+    metadata["nnCCDBInteractionRate"] = settings["nnCCDBInteractionRate"];
   }
   if (settings["nnCCDBBeamType"] != "") {
-    metadata[settings["nnCCDBBeamType"]] = settings["nnCCDBBeamType"];
+    metadata["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
+  bool retrieveSuccess = ccdbApi.retrieveBlob(settings["nnCCDBPath"], ".", metadata, 1, false, settings["outputFile"]);
+  // headers = ccdbApi.retrieveHeaders(settings["nnPathCCDB"], metadata, 1); // potentially needed to init some local variables
 
   if (retrieveSuccess) {
-    LOG(info) << "Network " << settings["nnPathCCDB"] << " retrieved from CCDB, stored at " << settings["networkPathLocal"];
+    LOG(info) << "Network " << settings["nnCCDBPath"] << " retrieved from CCDB, stored at " << settings["outputFile"];
   } else {
     LOG(error) << "Failed to retrieve network from CCDB";
   }
@@ -69,7 +70,7 @@ void GPUTPCNNClusterizerHost::init(const GPUSettingsProcessingNNclusterizer& set
 
   model_class.init(OrtOptions);
 
-  reg_model_paths = splitString(settings.nnRegressionPath, ":");
+  reg_model_paths = o2::utils::Str::tokenize(settings.nnRegressionPath, ':');
 
   if (!settings.nnClusterizerUseCfRegression) {
     if (model_class.getNumOutputNodes()[0][1] == 1 || reg_model_paths.size() == 1) {

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.h

@@ -37,17 +37,21 @@ class GPUTPCNNClusterizerHost
 {
  public:
   GPUTPCNNClusterizerHost() = default;
-  GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNclusterizer&, GPUTPCNNClusterizer&);
+  GPUTPCNNClusterizerHost(const GPUSettingsProcessingNNclusterizer&);
 
   void init(const GPUSettingsProcessingNNclusterizer&);
   void initClusterizer(const GPUSettingsProcessingNNclusterizer&, GPUTPCNNClusterizer&);
-  void loadFromCCDB(std::unordered_map<std::string, std::string>);
+  void loadFromCCDB(std::map<std::string, std::string>);
 
   void networkInference(o2::ml::OrtModel model, GPUTPCNNClusterizer& clusterer, size_t size, float* output, int32_t dtype);
 
   std::unordered_map<std::string, std::string> OrtOptions;
   o2::ml::OrtModel model_class, model_reg_1, model_reg_2; // For splitting clusters
   std::vector<std::string> reg_model_paths;
+
+  private:
+    std::map<std::string, std::string> metadata;
+    std::map<std::string, std::string> headers;
 }; // class GPUTPCNNClusterizerHost
 
 } // namespace o2::gpu