"docs/guides/vscode:/vscode.git/clone" did not exist on "699996e4fe1ebfd84c235fb721fbae2ed705016f"
Unverified Commit 70822f35 authored by Yan Ru Pei's avatar Yan Ru Pei Committed by GitHub
Browse files

fix: gate dynamo-memory numa module behind cfg(target_os = linux) (#6354)


Signed-off-by: default avatarPeaBrane <yanrpei@gmail.com>
Co-authored-by: default avatarCursor <cursoragent@cursor.com>
parent d38954c7
...@@ -14,6 +14,7 @@ ...@@ -14,6 +14,7 @@
pub mod actions; pub mod actions;
pub mod arena; pub mod arena;
pub mod nixl; pub mod nixl;
#[cfg(target_os = "linux")]
pub mod numa; pub mod numa;
/// Offset-based buffer views into underlying storage. /// Offset-based buffer views into underlying storage.
...@@ -41,6 +42,7 @@ pub use device::DeviceStorage; ...@@ -41,6 +42,7 @@ pub use device::DeviceStorage;
#[cfg(target_os = "linux")] #[cfg(target_os = "linux")]
pub use disk::DiskStorage; pub use disk::DiskStorage;
pub use external::ExternalDeviceMemory; pub use external::ExternalDeviceMemory;
#[cfg(target_os = "linux")]
pub use numa::{NumaNode, is_numa_enabled}; pub use numa::{NumaNode, is_numa_enabled};
pub use offset::OffsetBuffer; pub use offset::OffsetBuffer;
pub use pinned::PinnedStorage; pub use pinned::PinnedStorage;
......
...@@ -68,8 +68,6 @@ impl PinnedStorage { ...@@ -68,8 +68,6 @@ impl PinnedStorage {
/// - CUDA context creation fails /// - CUDA context creation fails
/// - Memory allocation fails /// - Memory allocation fails
pub fn new_for_device(len: usize, device_id: Option<u32>) -> Result<Self> { pub fn new_for_device(len: usize, device_id: Option<u32>) -> Result<Self> {
use super::numa;
if len == 0 { if len == 0 {
return Err(StorageError::AllocationFailed( return Err(StorageError::AllocationFailed(
"zero-sized allocations are not supported".into(), "zero-sized allocations are not supported".into(),
...@@ -80,24 +78,21 @@ impl PinnedStorage { ...@@ -80,24 +78,21 @@ impl PinnedStorage {
let ctx = cuda_context(gpu_id)?; let ctx = cuda_context(gpu_id)?;
let ptr = match device_id { let ptr = match device_id {
Some(gpu_id) if numa::is_numa_enabled() => { #[cfg(target_os = "linux")]
// NUMA-aware allocation via worker pool Some(gpu_id) if super::numa::is_numa_enabled() => {
tracing::debug!( tracing::debug!(
"Using NUMA-aware allocation for {} bytes on GPU {}", "Using NUMA-aware allocation for {} bytes on GPU {}",
len, len,
gpu_id gpu_id
); );
numa::worker_pool::NumaWorkerPool::global() super::numa::worker_pool::NumaWorkerPool::global()
.allocate_pinned_for_gpu(len, gpu_id) .allocate_pinned_for_gpu(len, gpu_id)
.map_err(StorageError::AllocationFailed)? as usize .map_err(StorageError::AllocationFailed)? as usize
} }
_ => { _ => unsafe {
// Direct allocation (no NUMA or device_id not specified)
unsafe {
ctx.bind_to_thread().map_err(StorageError::Cuda)?; ctx.bind_to_thread().map_err(StorageError::Cuda)?;
let ptr = let ptr = cudarc::driver::result::malloc_host(len, sys::CU_MEMHOSTALLOC_DEVICEMAP)
cudarc::driver::result::malloc_host(len, sys::CU_MEMHOSTALLOC_DEVICEMAP)
.map_err(StorageError::Cuda)?; .map_err(StorageError::Cuda)?;
let ptr = ptr as *mut u8; let ptr = ptr as *mut u8;
...@@ -106,8 +101,7 @@ impl PinnedStorage { ...@@ -106,8 +101,7 @@ impl PinnedStorage {
assert!(len < isize::MAX as usize); assert!(len < isize::MAX as usize);
ptr as usize ptr as usize
} },
}
}; };
Ok(Self { ptr, len, ctx }) Ok(Self { ptr, len, ctx })
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment