GPU: Clean up multi-threading, fix START/STOP/START for the future, don't poll for FMQ states, and better handling to avoid blocking FMQ state changes

Author: davidrohr

GPU/Workflow/include/GPUWorkflow/GPUWorkflowSpec.h

@@ -33,6 +33,7 @@ class TStopwatch;
 namespace fair::mq
 {
 struct RegionInfo;
+enum class State : int;
 } // namespace fair::mq
 namespace o2
 {
@@ -160,17 +161,18 @@ class GPURecoWorkflowSpec : public o2::framework::Task
   template <class D, class E, class F, class G, class H, class I, class J, class K>
   void processInputs(o2::framework::ProcessingContext&, D&, E&, F&, G&, bool&, H&, I&, J&, K&);
 
+  int runMain(o2::framework::ProcessingContext* pc, GPUTrackingInOutPointers* ptrs, GPUInterfaceOutputs* outputRegions, int threadIndex = 0, GPUInterfaceInputUpdate* inputUpdateCallback = nullptr);
   int runITSTracking(o2::framework::ProcessingContext& pc);
 
   int handlePipeline(o2::framework::ProcessingContext& pc, GPUTrackingInOutPointers& ptrs, gpurecoworkflow_internals::GPURecoWorkflowSpec_TPCZSBuffers& tpcZSmeta, o2::gpu::GPUTrackingInOutZS& tpcZS, std::unique_ptr<gpurecoworkflow_internals::GPURecoWorkflow_QueueObject>& context);
   void RunReceiveThread();
   void RunWorkerThread(int id);
-  void TerminateThreads();
+  void ExitPipeline();
   void handlePipelineEndOfStream(o2::framework::EndOfStreamContext& ec);
   void initPipeline(o2::framework::InitContext& ic);
   void enqueuePipelinedJob(GPUTrackingInOutPointers* ptrs, GPUInterfaceOutputs* outputRegions, gpurecoworkflow_internals::GPURecoWorkflow_QueueObject* context, bool inputFinal);
   void finalizeInputPipelinedJob(GPUTrackingInOutPointers* ptrs, GPUInterfaceOutputs* outputRegions, gpurecoworkflow_internals::GPURecoWorkflow_QueueObject* context);
-  int runMain(o2::framework::ProcessingContext* pc, GPUTrackingInOutPointers* ptrs, GPUInterfaceOutputs* outputRegions, int threadIndex = 0, GPUInterfaceInputUpdate* inputUpdateCallback = nullptr);
+  void receiveFMQStateCallback(fair::mq::State);
 
   CompletionPolicyData* mPolicyData;
   std::function<bool(o2::framework::DataProcessingHeader::StartTime)> mPolicyOrder;

GPU/Workflow/src/GPUWorkflowInternal.h

@@ -21,6 +21,7 @@
 #include <condition_variable>
 #include <queue>
 #include <array>
+#include <fairmq/States.h>
 
 namespace o2::gpu
 {
@@ -59,7 +60,12 @@ struct GPURecoWorkflow_QueueObject {
 struct GPURecoWorkflowSpec_PipelineInternals {
   std::mutex mutexDecodeInput;
 
-  fair::mq::Device* fmqDevice;
+  fair::mq::Device* fmqDevice = nullptr;
+
+  fair::mq::State fmqState = fair::mq::State::Undefined;
+  bool endOfStreamReceived = false;
+  std::mutex fmqStateMutex;
+  std::condition_variable fmqStateCheckNotify;
 
   std::thread receiveThread;
   std::mutex threadMutex;

GPU/Workflow/src/GPUWorkflowPipeline.cxx

@@ -34,6 +34,7 @@
 
 #include <fairmq/Device.h>
 #include <fairmq/Channel.h>
+#include <fairmq/States.h>
 
 using namespace o2::framework;
 using namespace o2::header;
@@ -45,6 +46,8 @@ using namespace o2::gpu::gpurecoworkflow_internals;
 namespace o2::gpu
 {
 
+static const std::string GPURecoWorkflowSpec_FMQCallbackKey = "GPURecoWorkflowSpec_FMQCallbackKey";
+
 struct pipelinePrepareMessage {
   static constexpr size_t MAGIC_WORD = 0X8473957353424134;
   size_t magicWord = MAGIC_WORD;
@@ -60,6 +63,7 @@ void GPURecoWorkflowSpec::initPipeline(o2::framework::InitContext& ic)
 {
   if (mSpecConfig.enableDoublePipeline == 1) {
     mPipeline->fmqDevice = ic.services().get<RawDeviceService>().device();
+    mPipeline->fmqDevice->SubscribeToStateChange(GPURecoWorkflowSpec_FMQCallbackKey, [this](fair::mq::State s) { receiveFMQStateCallback(s); });
     mPolicyOrder = [this](o2::framework::DataProcessingHeader::StartTime timeslice) {
       std::unique_lock lk(mPipeline->completionPolicyMutex);
       mPipeline->completionPolicyNotify.wait(lk, [pipeline = mPipeline.get()] { return pipeline->pipelineSenderTerminating || !pipeline->completionPolicyQueue.empty(); });
@@ -233,7 +237,11 @@ int GPURecoWorkflowSpec::handlePipeline(ProcessingContext& pc, GPUTrackingInOutP
 void GPURecoWorkflowSpec::handlePipelineEndOfStream(EndOfStreamContext& ec)
 {
   if (mSpecConfig.enableDoublePipeline == 1) {
-    TerminateThreads(); // TODO: Apparently breaks START / STOP / START
+    {
+      std::lock_guard lk(mPipeline->fmqStateMutex);
+      mPipeline->endOfStreamReceived = true;
+    }
+    mPipeline->fmqStateCheckNotify.notify_all();
   }
   if (mSpecConfig.enableDoublePipeline == 2) {
     auto* device = ec.services().get<RawDeviceService>().device();
@@ -247,18 +255,37 @@ void GPURecoWorkflowSpec::handlePipelineEndOfStream(EndOfStreamContext& ec)
   }
 }
 
+void GPURecoWorkflowSpec::receiveFMQStateCallback(fair::mq::State newState)
+{
+  {
+    std::lock_guard lk(mPipeline->fmqStateMutex);
+    if (mPipeline->fmqState != fair::mq::State::Running && newState == fair::mq::State::Running) {
+      mPipeline->endOfStreamReceived = false;
+    }
+    mPipeline->fmqState = newState;
+    if (newState == fair::mq::State::Exiting) {
+      mPipeline->fmqDevice->UnsubscribeFromStateChange(GPURecoWorkflowSpec_FMQCallbackKey);
+    }
+  }
+  mPipeline->fmqStateCheckNotify.notify_all();
+}
+
 void GPURecoWorkflowSpec::RunReceiveThread()
 {
   auto* device = mPipeline->fmqDevice;
-  while (device->GetCurrentState() != fair::mq::State::Running) {
-    usleep(300000);
-  }
   while (!mPipeline->shouldTerminate) {
     bool received = false;
     int recvTimeot = 1000;
     fair::mq::MessagePtr msg;
     LOG(debug) << "Waiting for out of band message";
     do {
+      {
+        std::unique_lock lk(mPipeline->fmqStateMutex);
+        mPipeline->fmqStateCheckNotify.wait(lk, [this]() { return (mPipeline->fmqState == fair::mq::State::Running && !mPipeline->endOfStreamReceived) || mPipeline->shouldTerminate; }); // Do not check mPipeline->fmqDevice->NewStatePending() since we wait for EndOfStream!
+      }
+      if (mPipeline->shouldTerminate) {
+        break;
+      }
       try {
         msg = device->NewMessageFor("gpu-prepare-channel", 0, 0);
         do {
@@ -279,8 +306,10 @@ void GPURecoWorkflowSpec::RunReceiveThread()
       LOG(fatal) << "Prepare message corrupted, invalid magic word";
     }
     if (m->flagEndOfStream) {
-      LOG(info) << "Received end-of-stream from out-of-band channel, terminating receive thread"; // TODO: Breaks START / STOP / START
-      break;
+      LOG(info) << "Received end-of-stream from out-of-band channel";
+      mPipeline->endOfStreamReceived = true;
+      mPipeline->mNTFReceived = 0;
+      continue;
     }
 
     auto context = std::make_unique<GPURecoWorkflow_QueueObject>();
@@ -317,7 +346,7 @@ void GPURecoWorkflowSpec::RunReceiveThread()
     context->ptrs.tpcZS = &context->tpcZS;
     context->ptrs.settingsTF = &context->tfSettings;
     context->mTFId = mPipeline->mNTFReceived;
-    if (mPipeline->mNTFReceived++ >= mPipeline->workers.size()) { // Do not inject the first 2 TF, since we need a first round of calib updates from DPL before starting
+    if (mPipeline->mNTFReceived++ >= mPipeline->workers.size()) { // Do not inject the first workers.size() TFs, since we need a first round of calib updates from DPL before starting
       enqueuePipelinedJob(&context->ptrs, nullptr, context.get(), false);
     }
     {
@@ -338,18 +367,22 @@ void GPURecoWorkflowSpec::RunReceiveThread()
   mPipeline->completionPolicyNotify.notify_one();
 }
 
-void GPURecoWorkflowSpec::TerminateThreads()
+void GPURecoWorkflowSpec::ExitPipeline()
 {
-  mPipeline->shouldTerminate = true;
-  for (unsigned int i = 0; i < mPipeline->workers.size(); i++) {
-    mPipeline->workers[i].inputQueueNotify.notify_one();
-  }
-  if (mPipeline->receiveThread.joinable()) {
-    mPipeline->receiveThread.join();
-  }
-  for (unsigned int i = 0; i < mPipeline->workers.size(); i++) {
-    if (mPipeline->workers[i].thread.joinable()) {
-      mPipeline->workers[i].thread.join();
+  if (mSpecConfig.enableDoublePipeline == 1 && mPipeline->fmqDevice) {
+    mPipeline->fmqDevice = nullptr;
+    mPipeline->shouldTerminate = true;
+    mPipeline->fmqStateCheckNotify.notify_all();
+    for (unsigned int i = 0; i < mPipeline->workers.size(); i++) {
+      mPipeline->workers[i].inputQueueNotify.notify_one();
+    }
+    if (mPipeline->receiveThread.joinable()) {
+      mPipeline->receiveThread.join();
+    }
+    for (unsigned int i = 0; i < mPipeline->workers.size(); i++) {
+      if (mPipeline->workers[i].thread.joinable()) {
+        mPipeline->workers[i].thread.join();
+      }
     }
   }
 }

GPU/Workflow/src/GPUWorkflowSpec.cxx

@@ -1211,7 +1211,7 @@ Outputs GPURecoWorkflowSpec::outputs()
 
 void GPURecoWorkflowSpec::deinitialize()
 {
-  TerminateThreads();
+  ExitPipeline();
   mQA.reset(nullptr);
   mDisplayFrontend.reset(nullptr);
   mGPUReco.reset(nullptr);