diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 0e9d21248..7a7d86dc4 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## 1.0.0 +### 1.0.0-rc.172 + +* Fixed `jfjoch_broker` cancelling every data collection with a CUDA "out of memory" error after long operation: GPU memory no longer leaks with each collection. + ### 1.0.0-rc.171 * Rugnux: basic support for CCD images (marCCD, SMV) and for gzipped miniCBF. diff --git a/image_analysis/indexing/CUDAMemHelpers.h b/image_analysis/indexing/CUDAMemHelpers.h index 40c6f0833..adbd98a20 100644 --- a/image_analysis/indexing/CUDAMemHelpers.h +++ b/image_analysis/indexing/CUDAMemHelpers.h @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include "../common/JFJochException.h" @@ -126,8 +127,8 @@ public: // the image loop, so do not hold everything. constexpr uint64_t CUDA_MEM_POOL_RELEASE_THRESHOLD = 1ull << 30; -// The stream device allocations are ordered on: one per (thread, device), created on first use and -// never destroyed. +// The stream device allocations are ordered on: one per (thread, device), borrowed on first use and +// handed back when the thread exits. The streams themselves are never destroyed. // // cudaMalloc and cudaFree are on CUDA's implicit-synchronization list - each one synchronises the // device across every stream - so one analysis engine per worker thread, each making a few dozen of @@ -137,24 +138,70 @@ constexpr uint64_t CUDA_MEM_POOL_RELEASE_THRESHOLD = 1ull << 30; // The stream is deliberately not owned by the engine: cudaFreeAsync must be ordered on a stream that // is still alive, and an engine's own stream can be destroyed before the buffers it allocated. It is // keyed by device because a thread pinned to one GPU must not order work on another's stream. +// +// Nor is it owned by the thread, for the same reason: a buffer can outlive the thread that allocated +// it - the shadow finder builds its accumulator on a throw-away thread and frees it from another, +// long after - so the thread only borrows the stream. But it does hand it back. The broker starts +// fresh worker threads for every data collection, and a stream holds about half a megabyte of device +// memory, so a stream per thread ever started is a leak that a long-running broker does not survive. +// Handed back, the streams number as many as the threads that were ever alive at once. +namespace jfjoch_cuda_allocation_streams { + struct Idle { + std::mutex m; + std::map> streams; // by device + }; + + // Never destroyed: a thread can still be handing a stream back while the process exits, and + // destroying streams during static destruction would call into a runtime that may be gone. + inline Idle &idle() { + static Idle *i = new Idle; + return *i; + } + + struct Borrowed { + std::map streams; // by device + ~Borrowed() { + std::lock_guard lock(idle().m); + for (auto &[device, stream] : streams) + idle().streams[device].push_back(std::move(stream)); + } + }; +} + inline cudaStream_t cuda_allocation_stream() { int device = 0; if (cudaGetDevice(&device) != cudaSuccess) return nullptr; - thread_local std::map streams; - const auto it = streams.find(device); - if (it != streams.end()) + thread_local jfjoch_cuda_allocation_streams::Borrowed borrowed; + const auto it = borrowed.streams.find(device); + if (it != borrowed.streams.end()) return it->second; - cudaStream_t stream = nullptr; - if (cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking) != cudaSuccess) + + auto &idle = jfjoch_cuda_allocation_streams::idle(); + { + std::lock_guard lock(idle.m); + auto &streams = idle.streams[device]; + if (!streams.empty()) { + CudaStream stream = std::move(streams.back()); + streams.pop_back(); + return borrowed.streams.emplace(device, std::move(stream)).first->second; + } + } + + try { + CudaStream stream; + cudaMemPool_t pool = nullptr; + if (cudaDeviceGetDefaultMemPool(&pool, device) == cudaSuccess) { + uint64_t threshold = CUDA_MEM_POOL_RELEASE_THRESHOLD; + cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold); + } + return borrowed.streams.emplace(device, std::move(stream)).first->second; + } catch (const JFJochException &) { + // No stream: the caller allocates synchronously instead, so this is handled - see + // CudaDevicePtr for why it must then not stay behind as the thread's last error. + cudaGetLastError(); return nullptr; - cudaMemPool_t pool = nullptr; - if (cudaDeviceGetDefaultMemPool(&pool, device) == cudaSuccess) { - uint64_t threshold = CUDA_MEM_POOL_RELEASE_THRESHOLD; - cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold); } - streams.emplace(device, stream); - return stream; } // Which allocator a device buffer uses. Pooled is what every per-worker buffer wants. Synchronous is @@ -183,14 +230,21 @@ public: explicit CudaDevicePtr(size_t count, CudaAlloc alloc = CudaAlloc::Pooled) { if (alloc == CudaAlloc::Pooled) stream_ = cuda_allocation_stream(); - if (stream_ && cudaMallocAsync(&ptr_, count * sizeof(T), stream_) == cudaSuccess) { - // The allocation is ordered on this stream, so anything that wants to use the memory from - // another stream has to be ordered after it. Waiting for it once here is that ordering, - // and it leaves the pointer as freely usable as cudaMalloc's would have been. - cudaStreamSynchronize(stream_); - return; + if (stream_) { + if (cudaMallocAsync(&ptr_, count * sizeof(T), stream_) == cudaSuccess) { + // The allocation is ordered on this stream, so anything that wants to use the memory + // from another stream has to be ordered after it. Waiting for it once here is that + // ordering, and it leaves the pointer as freely usable as cudaMalloc's would have been. + cudaStreamSynchronize(stream_); + return; + } + // The pool failing is handled - by the synchronous allocation below - so it must not stay + // behind as the thread's last error. Every kernel launch is followed by cudaGetLastError(), + // which would report it there: an "out of memory" thrown over an image whose buffers were + // all allocated and whose kernel launched. + cudaGetLastError(); + stream_ = nullptr; } - stream_ = nullptr; if (cudaMalloc(&ptr_, count * sizeof(T)) != cudaSuccess) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to allocate device memory"); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6d748a52b..45510abb0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -76,6 +76,7 @@ ADD_EXECUTABLE(jfjoch_test ImagePreprocessorGPUFusedTest.cpp AdaptiveSpotFinderGPUTest.cpp SpotExtractorGPUParityTest.cpp + CUDAMemHelpersTest.cpp CalcBraggPredictionTest.cpp SpotUtilsTest.cpp LatticeSearchTest.cpp diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp new file mode 100644 index 000000000..e9e00fc7c --- /dev/null +++ b/tests/CUDAMemHelpersTest.cpp @@ -0,0 +1,103 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include +#include "../common/CUDAWrapper.h" + +#ifdef JFJOCH_USE_CUDA + +#include +#include +#include +#include + +#include "../image_analysis/indexing/CUDAMemHelpers.h" + +// The broker starts fresh worker threads for every data collection. A stream per thread ever started +// is a device-memory leak that ends a long-running broker, so threads that follow one another have to +// end up on the same stream. +TEST_CASE("CudaAllocationStream_ReusedByLaterThreads", "[CUDAMemHelpers]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + std::set seen; + for (int i = 0; i < 100; i++) + std::thread([&seen] { + CudaDevicePtr buffer(1 << 16); + seen.insert(cuda_allocation_stream()); + }).join(); + + CHECK(seen.size() == 1); + CHECK(*seen.begin() != nullptr); +} + +TEST_CASE("CudaAllocationStream_DistinctForConcurrentThreads", "[CUDAMemHelpers]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + // Neither thread may exit - and hand its stream back - before the other has taken its own. + std::latch both_have_one(2); + auto take = [&both_have_one] { + const cudaStream_t stream = cuda_allocation_stream(); + both_have_one.count_down(); + both_have_one.wait(); + return stream; + }; + auto a = std::async(std::launch::async, take); + auto b = std::async(std::launch::async, take); + const cudaStream_t stream_a = a.get(); + const cudaStream_t stream_b = b.get(); + + CHECK(stream_a != nullptr); + CHECK(stream_b != nullptr); + CHECK(stream_a != stream_b); +} + +// The shadow finder builds its GPU accumulator on a throw-away thread and frees it from another one +// long after. The free is ordered on the allocating thread's stream, so that stream has to outlive +// the thread - including while a later thread has borrowed it. +TEST_CASE("CudaDevicePtr_FreedAfterAllocatingThreadExited", "[CUDAMemHelpers]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + auto buffer = std::async(std::launch::async, [] { + return std::make_unique>(1 << 20); + }).get(); + std::thread([] { CudaDevicePtr other(1 << 16); }).join(); + + REQUIRE(cudaMemset(buffer->get(), 0, 1 << 20) == cudaSuccess); + buffer.reset(); + CHECK(cudaDeviceSynchronize() == cudaSuccess); + CHECK(cudaGetLastError() == cudaSuccess); +} + +// A pool that cannot serve the request is not an error: the buffer comes from cudaMalloc instead. It +// must then not stay behind as the thread's last error, or the cudaGetLastError() after the next +// kernel launch throws "out of memory" over work that went fine. +TEST_CASE("CudaDevicePtr_HandledPoolFailureLeavesNoError", "[CUDAMemHelpers]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + int device = 0; + REQUIRE(cudaGetDevice(&device) == cudaSuccess); + + cudaMemPoolProps props{}; + props.allocType = cudaMemAllocationTypePinned; + props.location.type = cudaMemLocationTypeDevice; + props.location.id = device; + props.maxSize = 4 << 20; + + cudaMemPool_t previous = nullptr, capped = nullptr; + REQUIRE(cudaDeviceGetMemPool(&previous, device) == cudaSuccess); + REQUIRE(cudaMemPoolCreate(&capped, &props) == cudaSuccess); + REQUIRE(cudaDeviceSetMemPool(device, capped) == cudaSuccess); + { + CudaDevicePtr buffer(64 << 20); + CHECK(buffer.get() != nullptr); + CHECK(cudaGetLastError() == cudaSuccess); + } + REQUIRE(cudaDeviceSetMemPool(device, previous) == cudaSuccess); + cudaMemPoolDestroy(capped); +} + +#endif