From 060a89cd7d04361b49eab5bf9d6c9a28e5fe2ee4 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 19 Sep 2026 10:39:39 +0200 Subject: [PATCH] CUDA: allocation streams are handed back, not leaked; a handled pool failure leaves no error A long-running broker began cancelling every data collection with Device decoding failed (CUDA (GPU) error (out of memory)), falling back to host decompression CUDA (GPU) error (out of memory) while nvidia-smi showed the cards less than a fifth full. Two defects, both in CUDAMemHelpers.h and both from the pooled allocator (cudaMallocAsync) that came with rc.162. The leak. cuda_allocation_stream() kept one stream per (thread, device) in a thread_local map of raw cudaStream_t and never destroyed them. That was written against rugnux, where the worker threads live as long as the process. The broker starts fresh std::async threads for every data collection - 16 or 64 of them - so every collection left that many streams behind. Measured: 0.56 MB of device memory per leaked stream, linear to 4928 streams, never returned, with nothing on the host side growing. The RAII wrapper (CudaStream) was there but not used at this site, and using it as-is - a stream destroyed when its thread exits - would not have been safe: ShadowFinder builds its GPU accumulator on a throw-away std::async thread and frees it from another thread long after, and that free is ordered on the allocating thread's stream. So the streams are still never destroyed, but a thread now only borrows one: CudaStream objects live in a process-wide per-device idle list, a thread takes one on first use and hands it back when it exits. Their number is bounded by the threads that were ever alive at once instead of by the threads ever started. The list itself is deliberately leaked, so that nothing calls into CUDA during static destruction. Replaying the broker's pattern against the real header, 60 collections of 64 threads: before 3840 streams, 260 -> 2424 MB of device memory after 64 streams, 260 -> 358 MB The stale error. Every helper here throws a named message, yet the log carried the raw CUDA string, so the failure came through a cuda_err() and not from an allocation. CudaDevicePtr falls back to cudaMalloc when cudaMallocAsync fails, silently - but the failed call stays behind as the thread's last error, and the cudaGetLastError() that follows the next kernel launch reports it. The buffers were all allocated; the frame was lost anyway, once on the device-decode route (caught, hence the warning) and once more on the host fallback (fatal). The pooled attempt failing, and a stream that cannot be created, are both handled by falling back, so both now clear the error they leave. What finite resource the production cards ran out of at under 4 GB used was not established - no cap on the number of streams was found up to 4928 on the card this was measured on. The leak is the only thing on this path that grows with uptime. tests/CUDAMemHelpersTest.cpp: later threads end up on the same stream, concurrent threads on different ones, a buffer is freed cleanly after its allocating thread has exited (and another has borrowed its stream), and a pool that cannot serve a request leaves no error behind - the last by capping a memory pool at 4 MB so that the pooled attempt fails and the fallback succeeds. Co-Authored-By: Claude Fable 5.1 --- docs/CHANGELOG.md | 4 + image_analysis/indexing/CUDAMemHelpers.h | 94 ++++++++++++++++----- tests/CMakeLists.txt | 1 + tests/CUDAMemHelpersTest.cpp | 103 +++++++++++++++++++++++ 4 files changed, 182 insertions(+), 20 deletions(-) create mode 100644 tests/CUDAMemHelpersTest.cpp 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