mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 67 Wiki Activity

IndexerThreadPool: New version - requests are dropped, if no available workers (good for real-time, not good for offline processing...need smarter fix)

Browse Source
This commit is contained in:
leonarski_f
2026-04-16 22:35:32 +02:00
parent 6aa671c78c
commit 801a50830a
3 changed files with 181 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/indexing/IndexerThreadPool.cpp
+155 -117
View File
@@ -14,76 +14,24 @@
#include "FFTIndexerCPU.h"
#endif
IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, const NUMAHWPolicy &numa_policy)
: workers_ready(settings.GetIndexingThreads()),
viable_cell_min_spots(settings.GetViableCellMinSpots()),
stop(false) {
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();
}
}
IndexerResult IndexerThreadPool::Run(const DiffractionExperiment &experiment, const std::vector<Coord> &recip) {
if (experiment.GetIndexingAlgorithm() == IndexingAlgorithmEnum::None)
return IndexerResult{.lattice = {}, .indexing_time_s = 0};
if (recip.size() < viable_cell_min_spots)
return IndexerResult{.lattice = {}, .indexing_time_s = 0};
// 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.get();
}
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());
@@ -106,10 +54,20 @@ void IndexerThreadPool::Worker(int32_t threadIndex, const NUMAHWPolicy &numa_pol
}
} catch (const std::exception &e) {
spdlog::error("Failed to initialize GPU indexer: {}", e.what());
failed_start = true;
{
std::unique_lock<std::mutex> lock(m);
state = TaskState::ERROR;
}
c_running.notify_all();
return;
} catch (...) {
spdlog::error("Failed to initialize GPU indexer");
failed_start = true;
{
std::unique_lock<std::mutex> lock(m);
state = TaskState::ERROR;
}
c_running.notify_all();
return;
}
#endif
#ifdef JFJOCH_USE_FFTW
@@ -119,65 +77,145 @@ void IndexerThreadPool::Worker(int32_t threadIndex, const NUMAHWPolicy &numa_pol
fftw_indexer = std::make_unique<FFTIndexerCPU>(settings);
} catch (const std::exception &e) {
spdlog::error("Failed to initialize FFTW indexer: {}", e.what());
failed_start = true;
} catch (...) {
spdlog::error("Failed to initialize FFTW indexer");
failed_start = true;
}
#endif
workers_ready.count_down();
while (true) {
TaskPackage task;
{
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::milliseconds(50), [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
}
state = TaskState::ERROR;
}
try {
IndexerResult result;
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
{
std::unique_lock<std::mutex> lock(m);
state = TaskState::IDLE;
}
c_running.notify_all();
auto algorithm = task.experiment->GetIndexingAlgorithm();
Indexer *indexer = nullptr;
while (true) {
std::unique_ptr<TaskInput> input;
// Look for task + handle stop
{
std::unique_lock<std::mutex> lock(m);
c_start.wait(lock, [this] { return stop || state == TaskState::READY; });
if (stop)
return;
input = std::move(task_input);
}
if (input) {
std::unique_ptr<IndexerResult> tmp_result;
try {
auto algorithm = input->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();
} else if (algorithm == IndexingAlgorithmEnum::Auto) {
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"Internal error: Invalid indexing algorithm provided");
}
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);
c_done.notify_all();
}
// 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());
}
}
}
void IndexerThread::Finalize() {
{
std::unique_lock<std::mutex> lock(m);
stop = true;
}
c_start.notify_one();
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 std::move(tmp_result);
}
IndexerThread::~IndexerThread() {
Finalize();
}
IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings)
: worker_busy(settings.GetIndexingThreads(), 0),
viable_cell_min_spots(settings.GetViableCellMinSpots()) {
for (size_t i = 0; i < settings.GetIndexingThreads(); ++i)
tasks.emplace_back(std::make_unique<IndexerThread>(std::cref(settings), i));
}
IndexerResult IndexerThreadPool::Run(const DiffractionExperiment &experiment, const std::vector<Coord> &recip) {
if (experiment.GetIndexingAlgorithm() == IndexingAlgorithmEnum::None)
return IndexerResult{.lattice = {}, .indexing_time_s = 0};
if (recip.size() < viable_cell_min_spots)
return IndexerResult{.lattice = {}, .indexing_time_s = 0};
// Check if there is available worker
int task = -1;
{
std::unique_lock<std::mutex> lock(m);
for (int i = 0; i < tasks.size(); i++) {
if (worker_busy[i] == 0) {
task = i;
worker_busy[i] = 1;
break;
}
}
}
std::unique_ptr<IndexerResult> result;
if (task >= 0) {
result = tasks[task]->Run(experiment, recip);
{
std::unique_lock<std::mutex> lock(m);
worker_busy[task] = 0;
}
}
if (result)
return *result;
return IndexerResult{.lattice = {}, .indexing_time_s = 0};
}