dynamo_memory/numa/
worker_pool.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! NUMA worker pool for memory allocation with first-touch policy.
5//!
6//! This module provides dedicated worker threads that are pinned to specific NUMA nodes.
7//!
8//! ## Architecture
9//!
10//! - One worker thread per NUMA node (spawned lazily)
11//! - Workers pin themselves on startup (immune to application thread management)
12//! - Channel-based communication for allocation requests
13//! - First-touch page allocation ensures correct NUMA placement
14
15use super::get_current_cpu_numa_node;
16use cudarc::driver::result::malloc_host;
17use cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
18use nix::libc;
19use std::sync::mpsc::{Receiver, Sender, channel};
20use std::sync::{Arc, Mutex, OnceLock};
21use std::thread::{self, JoinHandle};
22use std::time::Duration;
23
24use super::{NumaNode, get_device_numa_node};
25
26/// Wrapper for raw pointer that can be sent between threads.
27///
28/// # Safety
29///
30/// This wrapper allows sending raw pointers across thread boundaries. The safety contract is:
31/// - The pointer is allocated by the worker thread and returned to the caller
32/// - The pointer is only dereferenced by the receiver (caller), never by the sender (worker)
33/// - Ownership is transferred: the caller is responsible for deallocation
34/// - The pointer remains valid for the lifetime expected by the caller
35struct SendPtr(*mut u8);
36
37// SAFETY: The pointer ownership is transferred from worker to caller.
38// The worker never accesses the pointer after sending it.
39unsafe impl Send for SendPtr {}
40
41/// Request to allocate CUDA pinned memory on a specific NUMA node.
42struct AllocRequest {
43    /// Number of bytes to allocate.
44    size: usize,
45    /// Target NUMA node for allocation.
46    node: NumaNode,
47    /// CUDA device ID (for context binding).
48    gpu_id: u32,
49    /// Channel for sending back the allocation result.
50    response: Sender<AllocResult>,
51}
52
53/// Result of allocation.
54type AllocResult = Result<SendPtr, String>;
55
56/// A dedicated worker thread pinned to a specific NUMA node.
57struct NumaWorker {
58    node: NumaNode,
59    request_tx: Option<Sender<AllocRequest>>,
60    handle: Option<JoinHandle<()>>,
61}
62
63impl NumaWorker {
64    /// Spawn a new worker thread pinned to the specified NUMA node.
65    fn spawn(node: NumaNode) -> Result<Self, String> {
66        let (request_tx, request_rx) = channel();
67
68        let handle = thread::Builder::new()
69            .name(format!("numa-worker-{}", node.0))
70            .spawn(move || {
71                Self::worker_loop(node, request_rx);
72            })
73            .map_err(|e| format!("Failed to spawn worker thread: {}", e))?;
74
75        Ok(Self {
76            node,
77            request_tx: Some(request_tx),
78            handle: Some(handle),
79        })
80    }
81
82    /// Worker thread main loop that processes allocation requests.
83    ///
84    /// On startup, the worker pins itself to the target NUMA node using
85    /// `sched_setaffinity`. It then processes allocation requests in a loop
86    /// until the channel is closed.
87    fn worker_loop(node: NumaNode, requests: Receiver<AllocRequest>) {
88        // First thing: pin this thread to the target NUMA node
89        tracing::trace!("Pinning worker thread to node {}", node.0);
90        if let Err(e) = super::pin_thread_to_numa_node(node) {
91            tracing::error!("Failed to pin worker thread to node {}: {}", node.0, e);
92            tracing::error!("Worker will continue but allocations may be suboptimal");
93        } else {
94            tracing::trace!("Successfully pinned worker thread to node {}", node.0);
95
96            // `pin_thread_to_numa_node` uses `sched_setaffinity` to set the CPU affinity mask
97            // but doesn't immediately migrate the thread. The scheduler will migrate at
98            // the next opportunity (timer tick, yield, etc).
99            // We yield once to give the scheduler a chance to migrate before we verify.
100            // This is primarily for accurate logging - allocations will happen on the right CPU
101            // regardless since the affinity mask prevents running on wrong CPUs.
102            thread::yield_now();
103            thread::sleep(Duration::from_millis(1));
104
105            // Verify we're on the right node
106            let current_node = super::get_current_cpu_numa_node();
107            tracing::trace!("Current node after pinning: {}", current_node.0);
108            if current_node != node {
109                tracing::warn!(
110                    "Worker thread on node {} after pinning (expected {})",
111                    current_node.0,
112                    node.0
113                );
114            } else {
115                tracing::trace!("NUMA worker thread for node {} started and pinned", node.0);
116            }
117        }
118
119        // Process allocation requests
120        loop {
121            tracing::trace!("Worker waiting for request on node {}", node.0);
122            match requests.recv() {
123                Ok(req) => {
124                    tracing::trace!(
125                        "Worker received CUDA pinned allocation request on node {}",
126                        node.0
127                    );
128                    let result = Self::do_cuda_pinned_allocation(req.size, req.node, req.gpu_id);
129                    match result {
130                        Ok(SendPtr(ptr)) => {
131                            if let Err(_e) = req.response.send(Ok(SendPtr(ptr))) {
132                                // Receiver gone: free to avoid leak
133                                tracing::warn!(
134                                    "Receiver dropped before receiving allocation, freeing {} bytes at {:p}",
135                                    req.size,
136                                    ptr
137                                );
138                                unsafe {
139                                    let _ = cudarc::driver::result::free_host(
140                                        ptr as *mut std::ffi::c_void,
141                                    );
142                                }
143                            }
144                        }
145                        Err(err) => {
146                            let _ = req.response.send(Err(err));
147                        }
148                    }
149                }
150                Err(_) => {
151                    // Channel closed, exit worker
152                    tracing::trace!(
153                        "NUMA worker for node {} shutting down (channel closed)",
154                        node.0
155                    );
156                    break;
157                }
158            }
159        }
160    }
161
162    /// Perform CUDA pinned memory allocation.
163    fn do_cuda_pinned_allocation(size: usize, node: NumaNode, gpu_id: u32) -> AllocResult {
164        if size == 0 {
165            return Err("Cannot allocate zero bytes".to_string());
166        }
167
168        // Verify we're on the correct NUMA node BEFORE allocation
169        let node_before = get_current_cpu_numa_node();
170        if node_before != node {
171            tracing::warn!(
172                "Worker thread moved! Expected NUMA node {}, currently on node {}",
173                node.0,
174                node_before.0
175            );
176        }
177
178        // Get or create CUDA context for this GPU
179        let ctx = crate::device::cuda_context(gpu_id)
180            .map_err(|e| format!("Failed to create CUDA context for device {}: {}", gpu_id, e))?;
181
182        unsafe {
183            // Bind CUDA context to this worker thread before allocation
184            // This ensures malloc_host has a valid context to work with
185            ctx.bind_to_thread()
186                .map_err(|e| format!("Failed to bind CUDA context to worker thread: {:?}", e))?;
187
188            // Verify thread is still on correct node after CUDA context binding
189            let node_after_ctx = get_current_cpu_numa_node();
190            if node_after_ctx != node {
191                tracing::warn!(
192                    "Thread moved after CUDA context bind! Expected node {}, now on node {}",
193                    node.0,
194                    node_after_ctx.0
195                );
196            }
197
198            // Allocate CUDA pinned memory
199            // This is called from the pinned worker thread, so pages will be
200            // allocated on the correct NUMA node via first-touch
201            let ptr = malloc_host(size, CU_MEMHOSTALLOC_DEVICEMAP)
202                .map_err(|e| format!("malloc_host failed: {:?}", e))?;
203
204            let ptr = ptr as *mut u8;
205
206            if ptr.is_null() {
207                return Err("malloc_host returned null".to_string());
208            }
209
210            // Verify thread is STILL on correct node before touching pages
211            let node_before_touch = get_current_cpu_numa_node();
212            if node_before_touch != node {
213                tracing::error!(
214                    "Thread on wrong node before first-touch! Expected {}, on node {} - memory will be misplaced!",
215                    node.0,
216                    node_before_touch.0
217                );
218            }
219
220            // Touch one byte per page to trigger first-touch policy efficiently
221            // This is much faster than zeroing the entire region for large allocations
222            let page_size = match libc::sysconf(libc::_SC_PAGESIZE) {
223                n if n > 0 => n as usize,
224                _ => 4096,
225            };
226            let mut offset = 0usize;
227            while offset < size {
228                std::ptr::write_volatile(ptr.add(offset), 0);
229                offset = offset.saturating_add(page_size);
230            }
231            // Ensure the last page is touched
232            if size > 0 && !size.is_multiple_of(page_size) {
233                std::ptr::write_volatile(ptr.add(size - 1), 0);
234            }
235
236            // Verify final node after touching
237            let node_after_touch = get_current_cpu_numa_node();
238
239            tracing::trace!(
240                "Worker allocated {} bytes (CUDA pinned) on GPU {} (target NUMA node {}) at {:p} - thread nodes: before={} after_ctx={} before_touch={} after_touch={}",
241                size,
242                gpu_id,
243                node.0,
244                ptr,
245                node_before.0,
246                node_after_ctx.0,
247                node_before_touch.0,
248                node_after_touch.0
249            );
250
251            Ok(SendPtr(ptr))
252        }
253    }
254
255    /// Request an allocation from this worker.
256    fn allocate(&self, size: usize, gpu_id: u32) -> AllocResult {
257        let (response_tx, response_rx) = channel();
258
259        let request = AllocRequest {
260            size,
261            node: self.node,
262            gpu_id,
263            response: response_tx,
264        };
265
266        self.request_tx
267            .as_ref()
268            .ok_or_else(|| "Worker has been shut down".to_string())?
269            .send(request)
270            .map_err(|_| "Worker thread has died".to_string())?;
271
272        // Wait for response with dynamic timeout based on allocation size
273        // Large allocations take time: we account for ~1 second per GB to touch pages
274        // Add 10 second base + 1 second per GB
275        let timeout_secs = 10u64 + (size as u64 / (1024 * 1024 * 1024));
276        let timeout = Duration::from_secs(timeout_secs.clamp(10, 300)); // Clamp to 10-300 seconds
277
278        tracing::trace!(
279            "Worker pool waiting for allocation of {} MB with timeout of {} seconds",
280            size / (1024 * 1024),
281            timeout.as_secs()
282        );
283
284        response_rx
285            .recv_timeout(timeout)
286            .map_err(|e| format!("Worker timeout after {} seconds: {}", timeout.as_secs(), e))?
287    }
288}
289
290impl Drop for NumaWorker {
291    fn drop(&mut self) {
292        tracing::trace!("Dropping NUMA worker for node {}", self.node.0);
293
294        // Drop request_tx FIRST to close the channel
295        // This causes recv() in worker thread to return Err and exit
296        self.request_tx.take();
297        tracing::trace!("Channel closed for worker node {}", self.node.0);
298
299        // Now the worker thread will exit its loop
300        if let Some(handle) = self.handle.take() {
301            tracing::trace!("Waiting for worker thread {} to join", self.node.0);
302            let _ = handle.join();
303            tracing::trace!("Worker thread {} joined", self.node.0);
304        }
305    }
306}
307
308/// Pool of NUMA workers, one per node.
309///
310/// This pool manages dedicated worker threads that are pinned to specific NUMA nodes.
311/// When you request an allocation for a GPU, the pool automatically determines the
312/// GPU's NUMA node and routes the request to the appropriate worker.
313pub struct NumaWorkerPool {
314    workers: Mutex<std::collections::HashMap<u32, Arc<NumaWorker>>>,
315}
316
317impl NumaWorkerPool {
318    fn new() -> Self {
319        Self {
320            workers: Mutex::new(std::collections::HashMap::new()),
321        }
322    }
323
324    /// Get the global worker pool.
325    ///
326    /// The pool is created lazily on first access and lives for the entire process lifetime.
327    pub fn global() -> &'static Self {
328        static POOL: OnceLock<NumaWorkerPool> = OnceLock::new();
329        POOL.get_or_init(NumaWorkerPool::new)
330    }
331
332    /// Get or create a worker for a NUMA node.
333    fn get_or_spawn_worker(&self, node: NumaNode) -> Result<Arc<NumaWorker>, String> {
334        let mut workers = self.workers.lock().unwrap();
335
336        if let Some(worker) = workers.get(&node.0) {
337            return Ok(worker.clone());
338        }
339
340        // Spawn new worker
341        let worker = NumaWorker::spawn(node)?;
342        let worker = Arc::new(worker);
343        workers.insert(node.0, worker.clone());
344
345        tracing::trace!("Spawned NUMA worker for node {}", node.0);
346
347        Ok(worker)
348    }
349
350    /// Allocate CUDA pinned memory for a specific GPU (auto-detects NUMA node).
351    ///
352    /// This method:
353    /// 1. Determines the GPU's NUMA node via CUDA driver PCI attributes + sysfs
354    /// 2. Routes the allocation to a worker pinned to that node
355    /// 3. The worker allocates and touches pages to ensure first-touch placement
356    ///
357    /// Returns `None` if the GPU's NUMA node cannot be determined, signaling
358    /// the caller to fall back to non-NUMA allocation.
359    ///
360    /// # Arguments
361    /// * `size` - Number of bytes to allocate
362    /// * `gpu_id` - CUDA device ID
363    ///
364    /// # Returns
365    /// `Some(ptr)` on success, `None` if NUMA node is unknown (caller should
366    /// use non-NUMA allocation). Returns `Err` on allocation failure.
367    pub fn allocate_pinned_for_gpu(
368        &self,
369        size: usize,
370        gpu_id: u32,
371    ) -> Result<Option<*mut u8>, String> {
372        let node = match get_device_numa_node(gpu_id) {
373            Some(node) => node,
374            None => {
375                tracing::debug!(
376                    "NUMA node unknown for GPU {}, skipping NUMA-aware allocation",
377                    gpu_id
378                );
379                return Ok(None);
380            }
381        };
382
383        tracing::debug!(
384            "Allocating {} bytes pinned memory for GPU {} (NUMA node {})",
385            size,
386            gpu_id,
387            node.0
388        );
389
390        let worker = self.get_or_spawn_worker(node)?;
391        worker
392            .allocate(size, gpu_id)
393            .map(|send_ptr| Some(send_ptr.0))
394    }
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400    use crate::numa::get_current_cpu_numa_node;
401
402    #[test]
403    fn test_worker_spawn() {
404        let node = NumaNode(0);
405        let worker = NumaWorker::spawn(node);
406        assert!(worker.is_ok());
407    }
408
409    #[test]
410    fn test_worker_pool_singleton() {
411        let pool1 = NumaWorkerPool::global();
412        let pool2 = NumaWorkerPool::global();
413        assert!(std::ptr::eq(pool1, pool2));
414    }
415
416    #[test]
417    fn test_get_current_cpu_numa_node() {
418        let node = get_current_cpu_numa_node();
419        if !node.is_unknown() {
420            println!("Current CPU on NUMA node: {}", node.0);
421        } else {
422            println!("NUMA node detection unavailable (single-node or fake NUMA)");
423        }
424    }
425
426    #[test]
427    fn test_numa_node_display() {
428        let node = NumaNode(0);
429        assert_eq!(format!("{}", node), "NumaNode(0)");
430
431        let unknown = NumaNode::UNKNOWN;
432        assert_eq!(format!("{}", unknown), "UNKNOWN");
433    }
434
435    #[test]
436    fn test_numa_node_is_unknown() {
437        let valid = NumaNode(0);
438        assert!(!valid.is_unknown());
439
440        let unknown = NumaNode::UNKNOWN;
441        assert!(unknown.is_unknown());
442    }
443}
444
445#[cfg(all(test, feature = "testing-cuda"))]
446mod cuda_tests {
447    use super::*;
448    use crate::numa::get_device_numa_node;
449
450    #[test]
451    fn test_worker_allocate_pinned() {
452        let node = NumaNode(0);
453        let worker = NumaWorker::spawn(node).unwrap();
454
455        let send_ptr = worker.allocate(4096, 0).unwrap();
456        let ptr = send_ptr.0;
457        assert!(!ptr.is_null());
458
459        unsafe {
460            cudarc::driver::result::free_host(ptr as *mut std::ffi::c_void).unwrap();
461        }
462    }
463
464    #[test]
465    fn test_worker_pool() {
466        let pool = NumaWorkerPool::new();
467
468        match pool.allocate_pinned_for_gpu(8192, 0).unwrap() {
469            Some(ptr) => unsafe {
470                assert!(!ptr.is_null());
471                cudarc::driver::result::free_host(ptr as *mut std::ffi::c_void).unwrap();
472            },
473            None => {
474                println!(
475                    "NUMA node unknown for GPU 0, allocation skipped (expected on single-socket)"
476                );
477            }
478        }
479    }
480
481    #[test]
482    fn test_worker_reuse() {
483        let pool = NumaWorkerPool::new();
484
485        // If NUMA node is unknown, both calls return None — that's fine
486        let r1 = pool.allocate_pinned_for_gpu(1024, 0).unwrap();
487        let r2 = pool.allocate_pinned_for_gpu(1024, 0).unwrap();
488
489        match (r1, r2) {
490            (Some(ptr1), Some(ptr2)) => unsafe {
491                assert!(!ptr1.is_null());
492                assert!(!ptr2.is_null());
493                assert_ne!(ptr1, ptr2);
494                cudarc::driver::result::free_host(ptr1 as *mut std::ffi::c_void).unwrap();
495                cudarc::driver::result::free_host(ptr2 as *mut std::ffi::c_void).unwrap();
496            },
497            (None, None) => {
498                println!("NUMA node unknown, both allocations skipped");
499            }
500            _ => panic!("inconsistent NUMA detection between two calls for same GPU"),
501        }
502    }
503
504    #[test]
505    fn test_zero_size_allocation_with_known_node() {
506        // Zero-size is rejected by the worker, but only if NUMA node is known.
507        // If NUMA node is unknown, allocate_pinned_for_gpu returns Ok(None) before
508        // reaching the worker.
509        let pool = NumaWorkerPool::new();
510        let result = pool.allocate_pinned_for_gpu(0, 0);
511        match result {
512            Ok(None) => {
513                println!("NUMA node unknown, zero-size check not reached");
514            }
515            Err(e) => {
516                assert!(e.contains("zero"));
517            }
518            Ok(Some(_)) => panic!("zero-size allocation should not succeed"),
519        }
520    }
521
522    #[test]
523    fn test_get_device_numa_node() {
524        let node = get_device_numa_node(0);
525        match node {
526            Some(n) => {
527                assert!(n.0 < 16, "NUMA node {} seems unreasonably high", n.0);
528                println!("GPU 0 on NUMA node: {}", n.0);
529            }
530            None => {
531                println!("GPU 0 has no determinable NUMA node");
532            }
533        }
534    }
535
536    #[test]
537    fn test_pinned_allocation_api() {
538        let pool = NumaWorkerPool::new();
539
540        if let Some(ptr) = pool.allocate_pinned_for_gpu(1024, 0).unwrap() {
541            assert!(!ptr.is_null());
542            unsafe {
543                cudarc::driver::result::free_host(ptr as *mut std::ffi::c_void).unwrap();
544            }
545        }
546    }
547
548    #[test]
549    fn test_worker_channel_communication() {
550        let node = NumaNode(0);
551        let worker = NumaWorker::spawn(node).unwrap();
552
553        let send_ptr = worker.allocate(1024, 0).unwrap();
554        let ptr = send_ptr.0;
555        assert!(!ptr.is_null());
556
557        unsafe {
558            cudarc::driver::result::free_host(ptr as *mut std::ffi::c_void).unwrap();
559        }
560    }
561}
dynamo_memory/numa/worker_pool.rs

dynamo_memory/numa/
worker_pool.rs
