Skip to content

Commit 9df0034

Browse files
authored
fix(cuda): bound concurrent stream-pool initialization (#10711)
## Problem Concurrent first use can create more CUDA streams than the pool's configured capacity. ## Solution Use `OnceCell::get_or_try_init` for each pool slot, replacing the custom initialization lock and `ArcSwap`. ## Assessment Initialized access remains lock-free and failed initialization can be retried. Uses `once_cell` because the fallible `std::sync::OnceLock` API is unstable. Based on #10726. Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>
1 parent e083440 commit 9df0034

3 files changed

Lines changed: 113 additions & 28 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎vortex-cuda/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ default = []
2121
_test-harness = []
2222

2323
[dependencies]
24-
arc-swap = { workspace = true }
2524
arrow-schema = { workspace = true, features = ["ffi"] }
2625
async-trait = { workspace = true }
2726
cudarc = { workspace = true, features = ["f16"] }
@@ -30,6 +29,7 @@ itertools = { workspace = true }
3029
kanal = { workspace = true }
3130
num-traits = { workspace = true }
3231
object_store = { workspace = true, features = ["fs"] }
32+
once_cell = { workspace = true }
3333
parking_lot = { workspace = true }
3434
prost = { workspace = true }
3535
tokio = { workspace = true, features = ["fs"] }

‎vortex-cuda/src/stream_pool.rs‎

Lines changed: 111 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -7,23 +7,22 @@ use std::sync::Arc;
77
use std::sync::atomic::AtomicUsize;
88
use std::sync::atomic::Ordering;
99

10-
use arc_swap::ArcSwapOption;
1110
use cudarc::driver::CudaContext;
1211
use cudarc::driver::CudaStream;
12+
use once_cell::sync::OnceCell;
1313
use vortex::error::VortexResult;
1414
use vortex::error::vortex_err;
1515

1616
use crate::stream::VortexCudaStream;
1717

1818
/// A pool of CUDA streams that hands out streams in a round-robin fashion.
1919
///
20-
/// Uses lock-free slot access via `ArcSwap` for high-concurrency scenarios.
21-
/// Streams are lazily created on first access to each slot and remain alive
22-
/// for the lifetime of the pool.
20+
/// Streams are lazily initialized once per slot and remain alive for the lifetime of the pool.
21+
/// Access to initialized slots is lock-free; failed initialization can be retried.
2322
pub struct VortexCudaStreamPool {
2423
context: Arc<CudaContext>,
2524
/// Fixed-size array of slots, each holding an optional stream.
26-
slots: Box<[ArcSwapOption<CudaStream>]>,
25+
slots: Box<[OnceCell<Arc<CudaStream>>]>,
2726
/// Round-robin counter for slot selection.
2827
next_index: AtomicUsize,
2928
}
@@ -46,7 +45,7 @@ impl VortexCudaStreamPool {
4645
/// * `max_capacity` - Maximum number of streams to maintain in the pool.
4746
pub fn new(context: Arc<CudaContext>, max_capacity: usize) -> Self {
4847
let slots = (0..max_capacity)
49-
.map(|_| ArcSwapOption::empty())
48+
.map(|_| OnceCell::new())
5049
.collect::<Vec<_>>()
5150
.into_boxed_slice();
5251

@@ -61,35 +60,23 @@ impl VortexCudaStreamPool {
6160
///
6261
/// Uses round-robin slot selection. If the selected slot has a stream,
6362
/// it is reused. Otherwise, a new stream is created for that slot.
64-
/// All operations are lock-free.
63+
/// Access to initialized slots is lock-free; initialization is serialized per slot.
6564
pub fn stream(&self) -> VortexResult<VortexCudaStream> {
6665
let slot_idx = self.next_index.fetch_add(1, Ordering::Relaxed) % self.slots.len();
67-
let slot = &self.slots[slot_idx];
66+
let stream = self.slots[slot_idx].get_or_try_init(|| {
67+
self.context
68+
.new_stream()
69+
.map_err(|e| vortex_err!("Failed to create CUDA stream: {}", e))
70+
})?;
6871

69-
// Fast path: stream already exists in slot.
70-
if let Some(stream) = slot.load_full() {
71-
return Ok(VortexCudaStream(stream));
72-
}
73-
74-
// Slow path: create a new stream.
75-
// Note: CudaContext::new_stream() already returns Arc<CudaStream>.
76-
let new_stream = self
77-
.context
78-
.new_stream()
79-
.map_err(|e| vortex_err!("Failed to create CUDA stream: {}", e))?;
80-
81-
// Store it in the slot. If another thread raced us, that's fine -
82-
// we'll just use our newly created stream this time.
83-
slot.store(Some(Arc::clone(&new_stream)));
84-
85-
Ok(VortexCudaStream(new_stream))
72+
Ok(VortexCudaStream(Arc::clone(stream)))
8673
}
8774

8875
/// Returns the current number of initialized streams in the pool.
8976
pub fn live_stream_count(&self) -> usize {
9077
self.slots
9178
.iter()
92-
.filter(|slot| slot.load().is_some())
79+
.filter(|slot| slot.get().is_some())
9380
.count()
9481
}
9582

@@ -98,3 +85,101 @@ impl VortexCudaStreamPool {
9885
self.slots.len()
9986
}
10087
}
88+
89+
#[cfg(test)]
90+
mod tests {
91+
use std::sync::Arc;
92+
use std::sync::Barrier;
93+
use std::thread;
94+
95+
use cudarc::driver::CudaContext;
96+
use vortex::error::VortexResult;
97+
use vortex::error::vortex_err;
98+
use vortex::utils::aliases::hash_set::HashSet;
99+
100+
use super::VortexCudaStreamPool;
101+
102+
#[crate::test]
103+
fn lazy_creation() -> VortexResult<()> {
104+
let context =
105+
CudaContext::new(0).map_err(|e| vortex_err!("Failed to create CUDA context: {e}"))?;
106+
let pool = VortexCudaStreamPool::new(context, 3);
107+
108+
assert_eq!(pool.max_capacity(), 3);
109+
assert_eq!(pool.live_stream_count(), 0);
110+
for expected_live in 1..=3 {
111+
let _stream = pool.stream()?;
112+
assert_eq!(pool.live_stream_count(), expected_live);
113+
}
114+
let _stream = pool.stream()?;
115+
assert_eq!(pool.live_stream_count(), 3);
116+
Ok(())
117+
}
118+
119+
#[crate::test]
120+
fn round_robin_reuse() -> VortexResult<()> {
121+
let context =
122+
CudaContext::new(0).map_err(|e| vortex_err!("Failed to create CUDA context: {e}"))?;
123+
let pool = VortexCudaStreamPool::new(context, 3);
124+
let streams = (0..pool.max_capacity())
125+
.map(|_| pool.stream())
126+
.collect::<VortexResult<Vec<_>>>()?;
127+
let handles = streams
128+
.iter()
129+
.map(|stream| stream.cu_stream() as usize)
130+
.collect::<HashSet<_>>();
131+
assert_eq!(handles.len(), pool.max_capacity());
132+
133+
for _ in 0..3 {
134+
for expected in &streams {
135+
let stream = pool.stream()?;
136+
assert!(Arc::ptr_eq(&stream.0, &expected.0));
137+
assert_eq!(stream.cu_stream(), expected.cu_stream());
138+
}
139+
}
140+
assert_eq!(pool.live_stream_count(), pool.max_capacity());
141+
Ok(())
142+
}
143+
144+
#[crate::test]
145+
fn concurrent_initialization_respects_capacity() -> VortexResult<()> {
146+
let context =
147+
CudaContext::new(0).map_err(|e| vortex_err!("Failed to create CUDA context: {e}"))?;
148+
for capacity in [1, 4] {
149+
let pool = VortexCudaStreamPool::new(Arc::clone(&context), capacity);
150+
let barrier = Barrier::new(32);
151+
let streams = thread::scope(|scope| {
152+
let workers = (0..32)
153+
.map(|_| {
154+
scope.spawn(|| {
155+
barrier.wait();
156+
(0..8)
157+
.map(|_| pool.stream())
158+
.collect::<VortexResult<Vec<_>>>()
159+
})
160+
})
161+
.collect::<Vec<_>>();
162+
workers
163+
.into_iter()
164+
.map(|worker| worker.join().expect("stream worker panicked"))
165+
.collect::<VortexResult<Vec<_>>>()
166+
})?;
167+
168+
// Keep every returned stream alive so transient handles cannot be recycled.
169+
let handles = streams
170+
.iter()
171+
.flatten()
172+
.map(|stream| stream.cu_stream() as usize)
173+
.collect::<HashSet<_>>();
174+
assert_eq!(handles.len(), capacity);
175+
assert_eq!(pool.live_stream_count(), capacity);
176+
let stored = (0..capacity)
177+
.map(|_| pool.stream())
178+
.collect::<VortexResult<Vec<_>>>()?;
179+
for stream in streams.iter().flatten() {
180+
assert!(stored.iter().any(|slot| Arc::ptr_eq(&stream.0, &slot.0)));
181+
}
182+
}
183+
Ok(())
184+
}
185+
}

0 commit comments

Comments
 (0)