mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 92 Wiki Activity
Files
rc184
Jungfraujoch/image_analysis/indexing/IndexerThreadPool.cpp
T
leonarski_fandClaude Opus 5.5 69ce5f3edd JFJochReceiverService::Indexing: kept FFT indexers of the old settings are released 
Since 0df766328 the broker keeps an FFTIndexerReuse for its whole life, so the broker's own indexing
pool (Lite and FPGA receivers) hands its GPU FFT indexers back to the reuse cache instead of freeing
them. The cache is keyed on the FFT settings and never shrank: every change of the indexing settings
through the API parked another 0.6-0.9 GB per worker on the card, until a pool build ran out of GPU
memory. (Before 0df766328 the same happened from the first live rotation run on.) The service now
releases the kept indexers when it rebuilds the pool; rugnux live keeps new ones at its next run.

(ported from rc184-2 8f47fb756)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQsU5a8AJUdJPgWevAK
2026-11-08 20:07:45 +01:00

510 lines
20 KiB
C++
Raw Permalink Blame History

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "IndexerThreadPool.h"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <map>
#include <tuple>
#include "../common/CUDAWrapper.h"
#include "../common/Logger.h"
#include "FFBIDXIndexerCPU.h"
#ifdef JFJOCH_USE_CUDA
#include "FFBIDXIndexer.h"
#include "FFTIndexerGPU.h"
#endif
#ifdef JFJOCH_USE_FFTW
#include "FFTIndexerCPU.h"
#endif
double WarmUpCuFFT() {
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() == 0)
return NAN;
// The first plan pays the initialisation, and compiles its kernels if the JIT cache does not have them
// yet; the next ones show what a plan costs from then on. The fastest of three, so that a stall
// beside (another card's context being made) is not taken for a slow plan.
double fastest_ms = INFINITY;
for (int i = 0; i < 4; i++) {
const auto t0 = std::chrono::steady_clock::now();
cufftHandle plan = 0;
if (cufftPlan3d(&plan, 64, 64, 64, CUFFT_R2C) != CUFFT_SUCCESS)
return NAN;
cufftDestroy(plan);
const double ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count();
if (i > 0)
fastest_ms = std::min(fastest_ms, ms);
}
return fastest_ms;
#else
return NAN;
#endif
}
#ifdef JFJOCH_USE_CUDA
FFTPlanReuse::FFTPlanReuse() : scope(std::make_unique<CudaFFTPlanScope>()) {}
#else
class CudaFFTPlanScope {};
FFTPlanReuse::FFTPlanReuse() = default;
#endif
FFTPlanReuse::~FFTPlanReuse() = default;
#ifndef JFJOCH_USE_CUDA
FFTIndexerReuse::FFTIndexerReuse() = default;
FFTIndexerReuse::~FFTIndexerReuse() = default;
void KeepFFTIndexers(const IndexingSettings &, int) {}
void ReleaseKeptFFTIndexers() {}
#endif
#ifdef JFJOCH_USE_CUDA
namespace {
// The card and what FFTIndexer's constructor reads from the settings.
using FFTIndexerKey = std::tuple<int, float, float, float, float, int64_t, float, int64_t>;
struct KeptFFTIndexers {
std::mutex m;
int scopes = 0;
std::map<FFTIndexerKey, std::vector<std::unique_ptr<Indexer>>> idle;
};
KeptFFTIndexers &kept_fft_indexers() {
static auto *k = new KeptFFTIndexers;
return *k;
}
FFTIndexerKey KeyOnCurrentCard(const IndexingSettings &s) {
int device = 0;
if (cudaGetDevice(&device) != cudaSuccess)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to read the CUDA device");
return {device, s.GetFFT_MaxUnitCell_A(), s.GetFFT_MinUnitCell_A(), s.GetFFT_MinAngle_deg(),
s.GetFFT_MaxAngle_deg(), s.GetFFT_NumVectors(), s.GetFFT_HighResolution_A(), s.GetRefineThreads()};
}
std::unique_ptr<Indexer> BorrowFFTIndexer(const IndexingSettings &settings) {
auto &k = kept_fft_indexers();
std::lock_guard lock(k.m);
auto it = k.idle.find(KeyOnCurrentCard(settings));
if (it == k.idle.end() || it->second.empty())
return nullptr;
auto ret = std::move(it->second.back());
it->second.pop_back();
return ret;
}
// On the card the indexer was built on, which is the calling worker's.
void GiveBackFFTIndexer(const IndexingSettings &settings, std::unique_ptr<Indexer> &indexer) {
if (!indexer)
return;
auto &k = kept_fft_indexers();
std::lock_guard lock(k.m);
if (k.scopes > 0)
k.idle[KeyOnCurrentCard(settings)].push_back(std::move(indexer));
}
}
FFTIndexerReuse::FFTIndexerReuse() {
std::lock_guard lock(kept_fft_indexers().m);
kept_fft_indexers().scopes++;
}
FFTIndexerReuse::~FFTIndexerReuse() {
std::map<FFTIndexerKey, std::vector<std::unique_ptr<Indexer>>> idle;
{
std::lock_guard lock(kept_fft_indexers().m);
if (--kept_fft_indexers().scopes == 0)
idle.swap(kept_fft_indexers().idle);
}
for (auto &[key, list] : idle) {
ScopedGpuPin pin(std::get<0>(key));
list.clear();
}
}
// An FFT indexer is a cuFFT plan and ~0.6-0.9 GB of device buffers, and several pipeline copies of one
// process run on the card at the same time: probe passes (GPUWorkBeside) beside the main line, for a
// few seconds each. So one that does not fit hands back what the allocation pool holds unused and waits
// for the work beside it to end, then is built again - as the merge's allocations do. The main line
// waits for the probes; a probe waits for the main line's image pass (GPUImagePass), never for a probe,
// itself included. Neither of those waits for the other, so nothing waits in a circle. Only when the
// indexer is built depends on this, never what it computes; one that does not fit the card fails as
// before.
static std::unique_ptr<Indexer> MakeFFTIndexerGPU(const IndexingSettings &settings, bool beside) {
if (auto kept = BorrowFFTIndexer(settings))
return kept;
try {
return std::make_unique<FFTIndexerGPU>(settings);
} catch (const JFJochException &) {
cuda_throw_if_context_lost(); // not memory: waiting would only hide the real error
cuda_clear_error();
}
if (beside)
wait_for_gpu_image_passes();
else
// That work is a probe pass of seconds, so this long means it is stuck: the indexer fails as
// it would have without it.
wait_for_gpu_work_beside(std::chrono::minutes(10));
cuda_trim_allocation_pool();
return std::make_unique<FFTIndexerGPU>(settings);
}
void ReleaseKeptFFTIndexers() {
std::map<FFTIndexerKey, std::vector<std::unique_ptr<Indexer>>> idle;
{
std::lock_guard lock(kept_fft_indexers().m);
idle.swap(kept_fft_indexers().idle);
}
for (auto &[key, list] : idle) {
ScopedGpuPin pin(std::get<0>(key));
list.clear();
}
}
void KeepFFTIndexers(const IndexingSettings &settings, int count) {
if (get_gpu_count() == 0)
return;
std::vector<std::unique_ptr<Indexer>> indexers;
for (int i = 0; i < count; i++)
indexers.push_back(MakeFFTIndexerGPU(settings, false));
for (auto &i : indexers)
GiveBackFFTIndexer(settings, i);
}
#endif
// The indexer for one RESOLVED algorithm, or nullptr if this build/host cannot serve it.
static std::unique_ptr<Indexer> MakeIndexer(IndexingAlgorithmEnum algorithm, const IndexingSettings &settings,
bool beside) {
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() > 0) {
if (algorithm == IndexingAlgorithmEnum::FFT)
return MakeFFTIndexerGPU(settings, beside);
if (algorithm == IndexingAlgorithmEnum::FFBIDX)
return std::make_unique<FFBIDXIndexer>();
}
#endif
if (algorithm == IndexingAlgorithmEnum::FFBIDX)
return std::make_unique<FFBIDXIndexerCPU>();
#ifdef JFJOCH_USE_FFTW
if (algorithm == IndexingAlgorithmEnum::FFTW)
return std::make_unique<FFTIndexerCPU>(settings);
#endif
return nullptr;
}
IndexerThread::IndexerThread(const IndexingSettings &settings, int threadid, IndexerConstruction construction,
bool beside)
: settings_(settings), construction_(construction), beside_(beside) {
std::unique_lock<std::mutex> lock(m);
state = TaskState::STARTING;
worker_thread = std::thread(&IndexerThread::Worker, this, 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");
}
}
void IndexerThread::Worker(int threadid) {
try {
pin_gpu();
} catch (const std::exception &e) {
spdlog::error("Failed to pin to GPU {}", e.what());
} catch (...) {
// GPU pinning errors are not critical and should be ignored for the time being.
}
std::unique_ptr<Indexer> fft_indexer, ffbidx_indexer, fftw_indexer;
// Preconstruct: build every indexer the requested algorithm could resolve to before the pool
// reports ready, so no cuFFT planning happens once frames are flowing, and a failure is fatal
// for the pool instead of being met frame by frame. OnFirstUse skips this and builds in the
// dispatch below.
if (construction_ == IndexerConstruction::Preconstruct) {
try {
const auto requested = settings_.GetAlgorithm();
if (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFT)
fft_indexer = MakeIndexer(IndexingAlgorithmEnum::FFT, settings_, beside_);
if (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFBIDX)
ffbidx_indexer = MakeIndexer(IndexingAlgorithmEnum::FFBIDX, settings_, beside_);
if ((requested == IndexingAlgorithmEnum::Auto && get_gpu_count() == 0)
|| requested == IndexingAlgorithmEnum::FFTW)
fftw_indexer = MakeIndexer(IndexingAlgorithmEnum::FFTW, settings_, beside_);
} catch (const std::exception &e) {
spdlog::error("Failed to initialize indexer: {}", e.what());
{
std::unique_lock<std::mutex> lock(m);
state = TaskState::ERROR;
}
c_running.notify_all();
return;
} catch (...) {
spdlog::error("Failed to initialize indexer");
{
std::unique_lock<std::mutex> lock(m);
state = TaskState::ERROR;
}
c_running.notify_all();
return;
}
}
{
std::unique_lock<std::mutex> lock(m);
state = TaskState::IDLE;
}
c_running.notify_all();
while (true) {
std::unique_ptr<TaskInput> input;
bool drop = false;
// Look for task + handle stop
{
std::unique_lock<std::mutex> lock(m);
c_start.wait(lock, [this] { return stop || drop_indexers || state == TaskState::READY; });
if (stop && (state != TaskState::READY)) {
#ifdef JFJOCH_USE_CUDA
GiveBackFFTIndexer(settings_, fft_indexer);
#endif
return;
}
if (state == TaskState::READY) {
state = TaskState::RUNNING;
input = std::move(task_input);
} else
drop = true;
}
if (drop) {
#ifdef JFJOCH_USE_CUDA
GiveBackFFTIndexer(settings_, fft_indexer);
#endif
fft_indexer.reset();
ffbidx_indexer.reset();
fftw_indexer.reset();
{
std::unique_lock<std::mutex> lock(m);
drop_indexers = false;
}
c_done.notify_all();
}
if (input) {
std::unique_ptr<IndexerResult> tmp_result;
try {
auto algorithm = input->experiment.GetIndexingAlgorithm();
std::unique_ptr<Indexer> *slot = nullptr;
switch (algorithm) {
case IndexingAlgorithmEnum::FFT: slot = &fft_indexer; break;
case IndexingAlgorithmEnum::FFBIDX: slot = &ffbidx_indexer; break;
case IndexingAlgorithmEnum::FFTW: slot = &fftw_indexer; break;
default: break;
}
// A preconstructing worker already holds it; an OnFirstUse worker builds it here,
// on the first frame that resolves to this algorithm.
if (slot && !*slot)
*slot = MakeIndexer(algorithm, settings_, beside_);
if (!slot || !*slot) {
// Algorithm is already resolved here (never Auto/None - see
// IndexerThreadPool::Run, which also checked this host can serve it). Reaching
// this means the resolved algorithm has no matching indexer in this build -
// fail loudly instead of silently not indexing.
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"Internal error: no indexer available for the resolved "
"indexing algorithm");
}
Indexer &indexer = **slot;
indexer.Setup(input->experiment);
tmp_result = std::make_unique<IndexerResult>(indexer.Run(input->recip, input->severity_only));
} catch (std::exception &e) {
// Hand the failure back as a result carrying the reason. A nullptr here was
// indistinguishable from a worker that was never dispatched, and both then read
// downstream as "this frame did not index".
spdlog::error("Indexer thread {} failed: {}", threadid, e.what());
tmp_result = std::make_unique<IndexerResult>(IndexerResult{
.lattice = {}, .indexing_time_s = 0, .executed = false,
.error = std::current_exception()});
// This thread goes on to the next frame, so a CUDA error left behind here would be
// reported by the first cudaGetLastError() of every frame that follows - one failed
// attempt would read as an indexer that never works again.
cuda_clear_error();
}
{
std::unique_lock<std::mutex> lock(m);
state = TaskState::COMPLETED;
result = std::move(tmp_result);
}
c_done.notify_all();
}
}
}
void IndexerThread::Release() {
std::unique_lock<std::mutex> lock(m);
if (stop)
return;
drop_indexers = true;
c_start.notify_all();
c_done.wait(lock, [this] { return !drop_indexers; });
}
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, bool severity_only) {
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), severity_only);
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, IndexerConstruction construction,
bool beside)
: worker_busy(settings.GetIndexingThreads(), 0),
worker_free_count(settings.GetIndexingThreads()),
viable_cell_min_spots(settings.GetViableCellMinSpots()),
blocking(settings.GetBlockingBehavior()),
settings_(settings) {
for (size_t i = 0; i < settings.GetIndexingThreads(); ++i)
tasks.emplace_back(std::make_unique<IndexerThread>(std::cref(settings), i, construction, beside));
}
void IndexerThreadPool::ReleaseIndexers() {
for (auto &t : tasks)
t->Release();
#ifdef JFJOCH_USE_CUDA
// Their buffers went back to the allocation pool, which would keep them reserved for the next
// buffer of whoever asks; hand them to the card.
cuda_trim_allocation_pools();
#endif
}
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,
bool severity_only) {
const auto algorithm = experiment.GetIndexingAlgorithm();
if (algorithm == IndexingAlgorithmEnum::None)
return IndexerResult{.lattice = {}, .indexing_time_s = 0, .executed = false};
// GetIndexingAlgorithm() must already have resolved Auto to a concrete algorithm;
// the pool has no policy to resolve it, so Auto here is an upstream contract bug.
if (algorithm == IndexingAlgorithmEnum::Auto)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"Internal error: indexing algorithm must be resolved (not Auto) "
"before reaching the indexer pool");
// The workers built their indexers from the raw requested algorithm, but the algorithm actually
// dispatched is the RESOLVED one (rotation, for instance, always resolves to the GPU FFT indexer
// when a GPU is present, ignoring the request). If the resolution lands on an algorithm this host
// did not build an indexer for, fail here with an explanation instead of the opaque "no indexer
// available for the resolved algorithm" from deep inside a worker.
const auto requested = experiment.GetIndexingSettings().GetAlgorithm();
const bool have_gpu = get_gpu_count() > 0;
#ifdef JFJOCH_USE_FFTW
constexpr bool fftw_built = true;
#else
constexpr bool fftw_built = false;
#endif
const bool servable =
(algorithm == IndexingAlgorithmEnum::FFT && have_gpu &&
(requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFT)) ||
(algorithm == IndexingAlgorithmEnum::FFBIDX &&
(requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFBIDX)) ||
(algorithm == IndexingAlgorithmEnum::FFTW && fftw_built &&
((requested == IndexingAlgorithmEnum::Auto && !have_gpu) || requested == IndexingAlgorithmEnum::FFTW));
if (!servable) {
std::string msg;
if (requested == IndexingAlgorithmEnum::FFTW && have_gpu)
msg = "FFTW is the CPU indexer and is not available on a node with a GPU. Rotation indexing "
"always uses the GPU FFT indexer here; select FFT or Auto, or run FFTW on a CPU-only node.";
else if (algorithm == IndexingAlgorithmEnum::FFT && !have_gpu)
msg = "FFT is the GPU indexer but no GPU is available. Select FFTW or Auto for CPU indexing.";
else if (algorithm == IndexingAlgorithmEnum::FFTW)
msg = "FFTW (CPU) indexing was requested but this build has no FFTW indexer.";
else
msg = "the requested indexing algorithm resolved to one with no indexer available on this host.";
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"Cannot index: " + msg);
}
// 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, severity_only);
} catch (const std::exception &e) {
spdlog::error("Indexer thread failed: {}", e.what());
result = std::make_unique<IndexerResult>(IndexerResult{
.lattice = {}, .indexing_time_s = 0, .executed = false,
.error = std::current_exception()});
}
{
std::unique_lock<std::mutex> lock(m);
worker_busy[task] = 0;
worker_free_count++;
}
c.notify_one();
}
if (result)
return *result;
// No free worker, or the pool is stopping: indexing was not attempted. Distinct from both a
// frame that did not index and an indexer that failed, and left without an error for that reason.
return IndexerResult{.lattice = {}, .indexing_time_s = 0};
}
Reference in new issue View Git Blame Copy Permalink