v1.0.0-rc.136 (#45)

This is an UNSTABLE release. The release has significant modifications and bug fixes, if things go wrong, it is better to revert to 1.0.0-rc.132.

* jfjoch_broker: Improve logic regarding indexing architecture and thread pools (work in progress).

Reviewed-on: #45

image_analysis/indexing/IndexerThreadPool.cpp

@@ -3,6 +3,7 @@
 
 #include "IndexerThreadPool.h"
 #include "../common/CUDAWrapper.h"
+#include "../common/Logger.h"
 
 #ifdef JFJOCH_USE_CUDA
 #include "FFBIDXIndexer.h"
@@ -13,71 +14,27 @@
 #include "FFTIndexerCPU.h"
 #endif
 
-IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, const NUMAHWPolicy &numa_policy)
-    : stop(false), workers_ready(settings.GetIndexingThreads()) {
-    for (size_t i = 0; i < settings.GetIndexingThreads(); ++i)
-        workers.emplace_back([this, i, numa_policy, settings] { Worker(i, numa_policy, settings); });
-    workers_ready.wait();
-
-    if (failed_start) {
-        {
-            std::unique_lock<std::mutex> lock(m);
-            stop = true;
-        }
-        cond.notify_all();
-
-        for (std::thread &worker: workers) {
-            if (worker.joinable())
-                worker.join();
-        }
-        throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
-                              "Cannot configure indexer (likely too many threads, not enough memory)");
+IndexerThread::IndexerThread(const IndexingSettings &settings, int threadid) {
+    std::unique_lock<std::mutex> lock(m);
+    state = TaskState::STARTING;
+    worker_thread = std::thread(&IndexerThread::Worker, this, std::cref(settings), threadid);
+    c_running.wait(lock, [this] { return state != TaskState::STARTING; });
+    if (state == TaskState::ERROR) {
+        worker_thread.join();
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
+                              "Indexer thread initialization failed");
     }
 }
 
-IndexerThreadPool::~IndexerThreadPool() { {
-        std::unique_lock<std::mutex> lock(m);
-        stop = true;
-    }
-    cond.notify_all();
-
-    for (std::thread &worker: workers) {
-        if (worker.joinable())
-            worker.join();
-    }
-}
-
-std::future<IndexerResult> IndexerThreadPool::Run(const DiffractionExperiment &experiment,
-                                                                   const std::vector<Coord>& recip) {
-    // Create a promise/future pair
-    auto promise = std::make_shared<std::promise<IndexerResult > >();
-    std::future<IndexerResult> result = promise->get_future(); {
-        std::unique_lock<std::mutex> lock(m);
-
-        // Don't allow enqueueing after stopping the pool
-        if (stop) {
-            throw std::runtime_error("Cannot enqueue on stopped thread pool");
-        }
-
-        // Create a task package with the data message and coordinates
-        taskQueue.emplace(TaskPackage{promise, &experiment, &recip});
-    }
-
-    cond.notify_one();
-    return result;
-}
-
-void IndexerThreadPool::Worker(int32_t threadIndex, const NUMAHWPolicy &numa_policy, const IndexingSettings &settings) {
+void IndexerThread::Worker(const IndexingSettings &settings, int threadid) {
     try {
 #ifdef JFJOCH_USE_CUDA
         auto gpu_count = get_gpu_count();
         if (gpu_count > 0)
-            NUMAHWPolicy::SelectGPUAndItsNUMA(threadIndex % gpu_count);
-        else
-            numa_policy.Bind(threadIndex);
-#else
-        numa_policy.Bind(threadIndex);
+            NUMAHWPolicy::SelectGPUAndItsNUMA(threadid % gpu_count);
 #endif
+    } catch (const std::exception &e) {
+        spdlog::error("Failed to bind thread to NUMA node: {}", e.what());
     } catch (...) {
         // NUMA policy errors are not critical and should be ignored for the time being.
     }
@@ -95,64 +52,186 @@ void IndexerThreadPool::Worker(int32_t threadIndex, const NUMAHWPolicy &numa_pol
                 || settings.GetAlgorithm() == IndexingAlgorithmEnum::FFBIDX)
                 ffbidx_indexer = std::make_unique<FFBIDXIndexer>();
         }
+    } catch (const std::exception &e) {
+        spdlog::error("Failed to initialize GPU indexer: {}", e.what());
+        {
+            std::unique_lock<std::mutex> lock(m);
+            state = TaskState::ERROR;
+        }
+        c_running.notify_all();
+        return;
     } catch (...) {
-        failed_start = true;
+        spdlog::error("Failed to initialize GPU indexer");
+        {
+            std::unique_lock<std::mutex> lock(m);
+            state = TaskState::ERROR;
+        }
+        c_running.notify_all();
+        return;
     }
 #endif
 #ifdef JFJOCH_USE_FFTW
-    if ((settings.GetAlgorithm() == IndexingAlgorithmEnum::Auto && (get_gpu_count() == 0))
-        || settings.GetAlgorithm() == IndexingAlgorithmEnum::FFTW)
-        fftw_indexer = std::make_unique<FFTIndexerCPU>(settings);
+    try {
+        if ((settings.GetAlgorithm() == IndexingAlgorithmEnum::Auto && (get_gpu_count() == 0))
+            || settings.GetAlgorithm() == IndexingAlgorithmEnum::FFTW)
+            fftw_indexer = std::make_unique<FFTIndexerCPU>(settings);
+    } catch (const std::exception &e) {
+        spdlog::error("Failed to initialize FFTW indexer: {}", e.what());
+        {
+            std::unique_lock<std::mutex> lock(m);
+            state = TaskState::ERROR;
+        }
+        c_running.notify_all();
+        return;
+    } catch (...) {
+        spdlog::error("Failed to initialize FFTW indexer");
+        {
+            std::unique_lock<std::mutex> lock(m);
+            state = TaskState::ERROR;
+        }
+        c_running.notify_all();
+        return;
+    }
 #endif
-
-    workers_ready.count_down();
+    {
+        std::unique_lock<std::mutex> lock(m);
+        state = TaskState::IDLE;
+    }
+    c_running.notify_all();
 
     while (true) {
-        TaskPackage task; {
+        std::unique_ptr<TaskInput> input;
+        // Look for task + handle stop
+        {
             std::unique_lock<std::mutex> lock(m);
-
-            // Add a timeout to the wait to ensure we can exit even if no notification
-            cond.wait_for(lock, std::chrono::seconds(1), [this] {
-                return stop || !taskQueue.empty();
-            });
-
-            // Check for exit conditions
-            if (stop && taskQueue.empty())
-                return; // Exit cleanly
-
-            if (!taskQueue.empty()) {
-                task = std::move(taskQueue.front());
-                taskQueue.pop();
-            } else {
-                continue; // No tasks, go back to waiting
-            }
+            c_start.wait(lock, [this] { return stop || state == TaskState::READY; });
+            if (stop && (state != TaskState::READY))
+                return;
+            state = TaskState::RUNNING;
+            input = std::move(task_input);
         }
-        try {
-            IndexerResult result;
+        if (input) {
+            std::unique_ptr<IndexerResult> tmp_result;
+            try {
+                auto algorithm = input->experiment.GetIndexingAlgorithm();
+                Indexer *indexer = nullptr;
 
-            auto algorithm = task.experiment->GetIndexingAlgorithm();
-            Indexer *indexer = nullptr;
+                if (algorithm == IndexingAlgorithmEnum::FFT && fft_indexer) {
+                    indexer = fft_indexer.get();
+                } else if (algorithm == IndexingAlgorithmEnum::FFBIDX && ffbidx_indexer) {
+                    indexer = ffbidx_indexer.get();
+                } else if (algorithm == IndexingAlgorithmEnum::FFTW && fftw_indexer) {
+                    indexer = fftw_indexer.get();
+                } else if (algorithm == IndexingAlgorithmEnum::Auto) {
+                    throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
+                                          "Internal error: Invalid indexing algorithm provided");
+                }
 
-            if (algorithm == IndexingAlgorithmEnum::FFT && fft_indexer) {
-                indexer = fft_indexer.get();
-            } else if (algorithm == IndexingAlgorithmEnum::FFBIDX && ffbidx_indexer) {
-                indexer = ffbidx_indexer.get();
-            } else if (algorithm == IndexingAlgorithmEnum::FFTW && fftw_indexer) {
-                indexer = fftw_indexer.get();
+                if (indexer) {
+                    indexer->Setup(input->experiment);
+                    tmp_result = std::make_unique<IndexerResult>(indexer->Run(input->recip));
+                }
+            } catch (std::exception &e) {
+                tmp_result = nullptr;
+                spdlog::error("Indexer thread {} failed: {}", threadid, e.what());
             }
-
-            if (indexer) {
-                indexer->Setup(*task.experiment);
-                result = indexer->Run(*task.recip);
+            {
+                std::unique_lock<std::mutex> lock(m);
+                state = TaskState::COMPLETED;
+                result = std::move(tmp_result);
             }
-
-            // Set the result via the promise
-            if (task.promise) {
-                task.promise->set_value(result);
-            }
-        } catch (std::exception &e) {
-            if (task.promise)
-                task.promise->set_exception(std::current_exception());
+            c_done.notify_all();
         }
     }
 }
+
+void IndexerThread::Finalize() {
+    {
+        std::unique_lock<std::mutex> lock(m);
+        stop = true;
+    }
+    c_start.notify_all();
+    if (worker_thread.joinable())
+        worker_thread.join();
+}
+
+std::unique_ptr<IndexerResult> IndexerThread::Run(const DiffractionExperiment &experiment,
+                                                  const std::vector<Coord> &recip) {
+    std::unique_ptr<IndexerResult> tmp_result;
+    {
+        std::unique_lock<std::mutex> lock(m);
+        if (stop)
+            return nullptr;
+        if (state != TaskState::IDLE)
+            return nullptr;
+        task_input = std::make_unique<TaskInput>(std::cref(experiment), std::cref(recip));
+        state = TaskState::READY;
+    }
+    c_start.notify_one();
+    {
+        std::unique_lock<std::mutex> lock(m);
+        c_done.wait(lock, [this] { return state == TaskState::COMPLETED; });
+        tmp_result = std::move(result);
+        state = TaskState::IDLE;
+    }
+    return tmp_result;
+}
+
+IndexerThread::~IndexerThread() {
+    Finalize();
+}
+
+IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings)
+    : worker_busy(settings.GetIndexingThreads(), 0),
+      worker_free_count(settings.GetIndexingThreads()),
+      viable_cell_min_spots(settings.GetViableCellMinSpots()),
+      blocking(settings.GetBlockingBehavior()) {
+    for (size_t i = 0; i < settings.GetIndexingThreads(); ++i)
+        tasks.emplace_back(std::make_unique<IndexerThread>(std::cref(settings), i));
+}
+
+int IndexerThreadPool::GetFreeWorker() {
+    std::unique_lock<std::mutex> lock(m);
+
+    if (tasks.size() == 0)
+        return -1;
+
+    if (blocking)
+        c.wait(lock, [this] { return worker_free_count > 0; });
+
+    for (int i = 0; i < tasks.size(); i++) {
+        if (worker_busy[i] == 0) {
+            worker_busy[i] = 1;
+            worker_free_count--;
+            return i;
+        }
+    }
+    return -1;
+}
+
+IndexerResult IndexerThreadPool::Run(const DiffractionExperiment &experiment, const std::vector<Coord> &recip) {
+    if (experiment.GetIndexingAlgorithm() == IndexingAlgorithmEnum::None)
+        return IndexerResult{.lattice = {}, .indexing_time_s = 0, .executed = false};
+
+    // Check if there is available worker
+   const int task = GetFreeWorker();
+
+    std::unique_ptr<IndexerResult> result;
+    if (task >= 0) {
+        try {
+            result = tasks[task]->Run(experiment, recip);
+        } catch (const std::exception &e) {
+            spdlog::error("Indexer thread failed: {}", e.what());
+            result = nullptr;
+        }
+        {
+            std::unique_lock<std::mutex> lock(m);
+            worker_busy[task] = 0;
+            worker_free_count++;
+        }
+        c.notify_one();
+    }
+    if (result)
+        return *result;
+    return IndexerResult{.lattice = {}, .indexing_time_s = 0};
+}