diff --git a/Cargo.lock b/Cargo.lock index ab869e730b9..33f17258b95 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9959,6 +9959,7 @@ dependencies = [ "parking_lot", "prost 0.14.4", "rstest", + "rustix", "tokio", "tracing", "vortex", diff --git a/Cargo.toml b/Cargo.toml index 7581d1c9653..c25713c0222 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -230,6 +230,7 @@ roaring = "0.11.0" rstest = "0.26.1" rstest_reuse = "0.7.0" rustc-hash = "2.1.1" +rustix = { version = "1.1", features = ["fs"] } serde = "1.0.220" serde_json = "1.0.138" serde_test = "1.0.176" diff --git a/vortex-cuda/Cargo.toml b/vortex-cuda/Cargo.toml index be6faa77e7b..a0c57478ce1 100644 --- a/vortex-cuda/Cargo.toml +++ b/vortex-cuda/Cargo.toml @@ -44,6 +44,9 @@ vortex-cuda-macros = { workspace = true } vortex-error = { workspace = true, features = ["object_store"] } vortex-nvcomp = { path = "nvcomp" } +[target.'cfg(target_os = "linux")'.dependencies] +rustix = { workspace = true } + [dev-dependencies] criterion = { workspace = true } futures = { workspace = true, features = ["executor"] } diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 1b07e12f0df..d83974921ec 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -23,3 +23,7 @@ Push host-resident arrays and close the sink using the standard `vx_array_sink_* Use `vx_cuda_scan_path_arrow_device_stream` to read such a local file through pinned host buffers and receive an Arrow C Device stream. Reuse the same CUDA session across scans so the pinned buffer pool and CUDA state are reused as well. + +On Linux, use `vx_cuda_scan_path_arrow_device_stream_with_options` with +`vx_cuda_scan_options.flags = VX_CUDA_SCAN_FLAG_DIRECT_IO` to bypass the operating system page +cache for pooled data-plane reads. Footer and zone-map reads remain buffered on the host. diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index a7c080e13cb..3c83b9b9ffd 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -94,6 +94,22 @@ vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session size_t block_rows, vx_error **error_out); +/** + * Options for scanning a CUDA-compatible Vortex file. + * + * Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. + */ +/** Bypass the operating system page cache for pooled data-plane reads. + * Footer and zone-map reads remain buffered. Supported only on Linux. */ +#define VX_CUDA_SCAN_FLAG_DIRECT_IO (UINT32_C(1) << 0) + +typedef struct vx_cuda_scan_options { + /** Bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. */ + uint32_t flags; + /** Number of rows in each output ArrowDeviceArray. Zero uses layout-derived splitting. */ + size_t batch_rows; +} vx_cuda_scan_options; + /** * Scan a local CUDA-compatible Vortex file as an Arrow C Device stream. * @@ -125,6 +141,19 @@ int vx_cuda_scan_path_arrow_device_stream_batch_rows(const vx_session *session, struct ArrowDeviceArrayStream *out_stream, vx_error **error_out); +/** + * Scan a local CUDA-compatible Vortex file with explicit options. + * + * This has the same ownership and file compatibility requirements as + * `vx_cuda_scan_path_arrow_device_stream`. Pass NULL or a zero-initialized options struct to use + * buffered file I/O and layout-derived batch splitting. + */ +int vx_cuda_scan_path_arrow_device_stream_with_options(const vx_session *session, + vx_view path, + const vx_cuda_scan_options *options, + struct ArrowDeviceArrayStream *out_stream, + vx_error **error_out); + /** * Export a borrowed Vortex array for cuDF's Arrow Device import path. * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index be2a9defc59..623e017d05d 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -25,6 +25,7 @@ use vortex::session::SessionExt; use vortex::session::VortexSession; use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; +use vortex_cuda::PooledFileReadAtOptions; use vortex_cuda::arrow::ArrowDeviceArray; use vortex_cuda::arrow::ArrowDeviceArrayStream; use vortex_cuda::arrow::DeviceArrayExt; @@ -49,6 +50,22 @@ use vortex_ffi::vx_view; const VX_CUDA_OK: c_int = 0; const VX_CUDA_ERR: c_int = 1; +/// Enable direct I/O for pooled CUDA file reads. +pub const VX_CUDA_SCAN_FLAG_DIRECT_IO: u32 = 1 << 0; +const VX_CUDA_SCAN_KNOWN_FLAGS: u32 = VX_CUDA_SCAN_FLAG_DIRECT_IO; + +/// Options for scanning a CUDA-compatible Vortex file. +/// +/// Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. +#[repr(C)] +#[derive(Default)] +pub struct vx_cuda_scan_options { + /// A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. + pub flags: u32, + /// Number of rows in each output batch. Zero uses layout-derived splitting. + pub batch_rows: usize, +} + /// Return a Vortex session with a [`CudaSession`] session variable. /// /// If `session` already has CUDA support, this returns a clone of it. Otherwise it @@ -142,7 +159,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( }) } -/// Scan a local Vortex file and export an Arrow C Device stream. +/// Scan a local Vortex file with buffered I/O and export an Arrow C Device stream. /// /// Footer and zone-map reads remain on the host. Data segments are staged through pinned host /// buffers and transferred directly to the GPU. @@ -171,7 +188,13 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( error_out: *mut *mut vx_error, ) -> c_int { unsafe { - vx_cuda_scan_path_arrow_device_stream_batch_rows(session, path, 0, out_stream, error_out) + vx_cuda_scan_path_arrow_device_stream_with_options( + session, + path, + ptr::null(), + out_stream, + error_out, + ) } } @@ -195,19 +218,60 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_batch_rows batch_rows: usize, out_stream: *mut ArrowDeviceArrayStream, error_out: *mut *mut vx_error, +) -> c_int { + let options = vx_cuda_scan_options { + batch_rows, + ..Default::default() + }; + unsafe { + vx_cuda_scan_path_arrow_device_stream_with_options( + session, + path, + &raw const options, + out_stream, + error_out, + ) + } +} + +/// Scan a local Vortex file with explicit options and export an Arrow C Device stream. +/// +/// This has the same ownership and file compatibility requirements as +/// [`vx_cuda_scan_path_arrow_device_stream`]. Pass a null `options` pointer or a zero-initialized +/// [`vx_cuda_scan_options`] to use buffered file I/O and layout-derived batch splitting. +/// +/// # Safety +/// +/// `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the +/// duration of this call and contain UTF-8. `options`, when non-null, must point to a valid +/// [`vx_cuda_scan_options`]. `out_stream` must be a valid writable pointer. If `error_out` is +/// non-null, it must be valid for writing one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_options( + session: *const vx_session, + path: vx_view, + options: *const vx_cuda_scan_options, + out_stream: *mut ArrowDeviceArrayStream, + error_out: *mut *mut vx_error, ) -> c_int { try_or(error_out, VX_CUDA_ERR, || { vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output"); let path = unsafe { path.as_str() }?.to_owned(); let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?; + let options = unsafe { scan_options(options) }?; let array_stream = ffi_runtime().block_on(async { - let file = session.open_options().with_cuda().open_path(path).await?; + let file = session + .open_options() + .with_cuda() + .with_read_at_options(options.read_at_options) + .open_path(path) + .await?; let scan = file.scan()?; - let scan = if batch_rows == 0 { + let scan = if options.batch_rows == 0 { scan } else { - scan.with_split_by(SplitBy::RowCount(batch_rows)) + scan.with_split_by(SplitBy::RowCount(options.batch_rows)) }; Ok::<_, vortex::error::VortexError>(scan.into_array_stream()?.boxed()) })?; @@ -218,6 +282,46 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_batch_rows }) } +struct CudaScanOptions { + read_at_options: PooledFileReadAtOptions, + batch_rows: usize, +} + +unsafe fn scan_options(options: *const vx_cuda_scan_options) -> VortexResult { + let (flags, batch_rows) = if options.is_null() { + (0, 0) + } else { + let options = unsafe { &*options }; + (options.flags, options.batch_rows) + }; + vortex_ensure!( + flags & !VX_CUDA_SCAN_KNOWN_FLAGS == 0, + "unsupported CUDA scan option flags: {:#x}", + flags & !VX_CUDA_SCAN_KNOWN_FLAGS + ); + + let read_at_options = PooledFileReadAtOptions::default(); + let read_at_options = if flags & VX_CUDA_SCAN_FLAG_DIRECT_IO == 0 { + read_at_options + } else { + #[cfg(target_os = "linux")] + { + read_at_options.with_direct_io() + } + #[cfg(not(target_os = "linux"))] + { + return Err(vortex::error::vortex_err!( + "direct CUDA file I/O is only supported on Linux" + )); + } + }; + + Ok(CudaScanOptions { + read_at_options, + batch_rows, + }) +} + /// Export a borrowed Vortex array for cuDF's Arrow Device import path. /// /// On success returns `0` and writes independently releasable `out_schema` and `out_array`; the @@ -321,6 +425,49 @@ mod tests { use super::*; + #[test] + fn scan_options_default_to_buffered_io() { + let options = vx_cuda_scan_options::default(); + assert_eq!(options.flags, 0); + assert_eq!(options.batch_rows, 0); + } + + #[test] + fn rejects_unknown_scan_option_flags() { + let options = vx_cuda_scan_options { + flags: 1 << 31, + ..Default::default() + }; + assert!(unsafe { scan_options(&raw const options) }.is_err()); + } + + #[cfg(target_os = "linux")] + #[test] + fn maps_direct_io_scan_option_to_pooled_reader() -> VortexResult<()> { + let options = vx_cuda_scan_options { + flags: VX_CUDA_SCAN_FLAG_DIRECT_IO, + ..Default::default() + }; + assert_eq!( + unsafe { scan_options(&raw const options) }?.read_at_options, + PooledFileReadAtOptions::default().with_direct_io() + ); + Ok(()) + } + + #[test] + fn maps_batch_rows_scan_option() -> VortexResult<()> { + let options = vx_cuda_scan_options { + batch_rows: 8192, + ..Default::default() + }; + assert_eq!( + unsafe { scan_options(&raw const options) }?.batch_rows, + 8192 + ); + Ok(()) + } + fn test_session(session: VortexSession) -> *mut vx_session { Box::into_raw(Box::new(session)).cast::() } diff --git a/vortex-cuda/src/file.rs b/vortex-cuda/src/file.rs index 9a41f51f9b9..932c390b0ff 100644 --- a/vortex-cuda/src/file.rs +++ b/vortex-cuda/src/file.rs @@ -20,6 +20,7 @@ use vortex::layout::segments::SegmentSource; use crate::CudaSessionExt; use crate::PooledFileReadAt; +use crate::PooledFileReadAtOptions; use crate::layout::register_cuda_layout; /// Extension trait for opening CUDA-readable files from [`VortexOpenOptions`]. @@ -31,7 +32,10 @@ pub trait CudaOpenOptionsExt { impl CudaOpenOptionsExt for VortexOpenOptions { fn with_cuda(self) -> CudaOpenOptions { - CudaOpenOptions { inner: self } + CudaOpenOptions { + inner: self, + read_at_options: PooledFileReadAtOptions::default(), + } } } @@ -41,9 +45,18 @@ impl CudaOpenOptionsExt for VortexOpenOptions { /// configured before calling `with_cuda`. pub struct CudaOpenOptions { inner: VortexOpenOptions, + read_at_options: PooledFileReadAtOptions, } impl CudaOpenOptions { + /// Configure how the pooled data-plane reader opens and reads the local file. + /// + /// Footer and zone-map reads continue to use the standard buffered host reader. + pub fn with_read_at_options(mut self, options: PooledFileReadAtOptions) -> Self { + self.read_at_options = options; + self + } + /// Open a local Vortex file for CUDA execution. /// /// The footer and zone-map segments are read through the ordinary host path. All other file @@ -65,11 +78,12 @@ impl CudaOpenOptions { let pool = Arc::clone(cuda_session.pinned_buffer_pool()); drop(cuda_session); - let reader = Arc::new(PooledFileReadAt::open( + let reader = Arc::new(PooledFileReadAt::open_with_options( &path, session.handle(), pool, stream, + self.read_at_options, )?); let data_file = data_options .with_footer(footer.clone()) diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 69a2f7263eb..3c712d20fb8 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -62,6 +62,7 @@ pub use pinned::PinnedPoolStats; pub use pinned::PooledPinnedBuffer; pub use pooled_read_at::PooledByteBufferReadAt; pub use pooled_read_at::PooledFileReadAt; +pub use pooled_read_at::PooledFileReadAtOptions; pub use pooled_read_at::PooledObjectStoreReadAt; pub use session::CudaSession; pub use session::CudaSessionExt; diff --git a/vortex-cuda/src/pinned.rs b/vortex-cuda/src/pinned.rs index 52b1aa46e2a..11796c5899a 100644 --- a/vortex-cuda/src/pinned.rs +++ b/vortex-cuda/src/pinned.rs @@ -309,6 +309,19 @@ impl PooledPinnedBuffer { .unwrap_or_else(|e| vortex_panic!("failed to access pinned host buffer: {e}")) } + /// Shortens the logical length without reallocating the pinned buffer. + /// + /// This is an O(1) metadata update; the allocation capacity is unchanged. + #[cfg(target_os = "linux")] + pub(crate) fn truncate(&mut self, len: usize) { + let inner = self + .inner + .as_mut() + .unwrap_or_else(|| vortex_panic!("buffer already consumed")); + assert!(len <= inner.len()); + inner.set_logical_len(len); + } + /// Submits a non-blocking H2D DMA transfer and returns a device buffer. /// /// The pinned buffer is placed in the pool's inflight queue, gated on a `CudaEvent` marking diff --git a/vortex-cuda/src/pooled_read_at.rs b/vortex-cuda/src/pooled_read_at.rs index d3024140c96..f5352a7c06a 100644 --- a/vortex-cuda/src/pooled_read_at.rs +++ b/vortex-cuda/src/pooled_read_at.rs @@ -1,9 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fs::File; use std::io; -use std::path::Path; use std::sync::Arc; use futures::FutureExt; @@ -30,98 +28,13 @@ use vortex::io::std_file::read_exact_at; use crate::pinned::PinnedByteBufferPool; use crate::stream::VortexCudaStream; -/// Default number of concurrent requests to allow for local file I/O. -pub const DEFAULT_FILE_CONCURRENCY: usize = 32; -/// Default number of concurrent requests to allow for object store I/O. -pub const DEFAULT_OBJECT_STORE_CONCURRENCY: usize = 192; - -/// File reader that uses CUDA pinned host memory for I/O buffers and transfers -/// directly to the GPU. -/// -/// Reads into a pooled pinned (page-locked) buffer, then submits a non-blocking -/// H2D DMA transfer and returns a device `BufferHandle`. -/// -/// This is a data-plane reader. To open a complete local Vortex file, prefer -/// [`crate::CudaOpenOptionsExt::with_cuda`], which keeps the footer and zone maps on the host. -#[derive(Clone)] -pub struct PooledFileReadAt { - uri: Arc, - file: Arc, - handle: Handle, - pool: Arc, - stream: VortexCudaStream, -} - -impl PooledFileReadAt { - /// Open a file for pooled reading with direct device transfer. - pub fn open( - path: impl AsRef, - handle: Handle, - pool: Arc, - stream: VortexCudaStream, - ) -> VortexResult { - let path = path.as_ref(); - let uri = Arc::from(path.to_string_lossy().to_string()); - let file = Arc::new(File::open(path)?); - Ok(Self { - uri, - file, - handle, - pool, - stream, - }) - } -} - -impl VortexReadAt for PooledFileReadAt { - fn uri(&self) -> Option<&Arc> { - Some(&self.uri) - } - - fn coalesce_config(&self) -> Option { - Some(CoalesceConfig::file()) - } +mod file; - fn concurrency(&self) -> usize { - DEFAULT_FILE_CONCURRENCY - } - - fn size(&self) -> BoxFuture<'static, VortexResult> { - let file = Arc::clone(&self.file); - async move { - let metadata = file.metadata()?; - Ok(metadata.len()) - } - .boxed() - } - - fn read_at( - &self, - offset: u64, - length: usize, - _alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { - let file = Arc::clone(&self.file); - let handle = self.handle.clone(); - let stream = self.stream.clone(); - let pool = Arc::clone(&self.pool); - - async move { - let mut target = pool.get(length)?; - let target = handle - .spawn_blocking(move || { - read_exact_at(&file, target.as_mut_slice(), offset)?; - Ok::<_, io::Error>(target) - }) - .await - .map_err(VortexError::from)?; +pub use file::PooledFileReadAt; +pub use file::PooledFileReadAtOptions; - let cuda_buf = target.transfer_to_device(&stream)?; - Ok(BufferHandle::new_device(Arc::new(cuda_buf))) - } - .boxed() - } -} +/// Default number of concurrent requests to allow for object store I/O. +pub const DEFAULT_OBJECT_STORE_CONCURRENCY: usize = 192; /// Object store reader that uses CUDA pinned host memory for I/O buffers and /// transfers directly to the GPU. diff --git a/vortex-cuda/src/pooled_read_at/file.rs b/vortex-cuda/src/pooled_read_at/file.rs new file mode 100644 index 00000000000..b2b44fed2bf --- /dev/null +++ b/vortex-cuda/src/pooled_read_at/file.rs @@ -0,0 +1,231 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(target_os = "linux")] +mod direct; + +use std::fs::File; +use std::ops::Range; +use std::path::Path; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::BoxFuture; +use vortex::array::buffer::BufferHandle; +use vortex::buffer::Alignment; +use vortex::error::VortexResult; +use vortex::io::CoalesceConfig; +use vortex::io::VortexReadAt; +use vortex::io::runtime::Handle; +use vortex::io::std_file::read_exact_at; + +#[cfg(target_os = "linux")] +use self::direct::DirectFileReadBackend; +use crate::pinned::PinnedByteBufferPool; +use crate::pinned::PooledPinnedBuffer; +use crate::stream::VortexCudaStream; + +/// Default number of concurrent requests to allow for local file I/O. +pub const DEFAULT_FILE_CONCURRENCY: usize = 32; + +/// Options controlling how [`PooledFileReadAt`] opens and reads a local file. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct PooledFileReadAtOptions { + direct_io: bool, +} + +impl PooledFileReadAtOptions { + /// Bypass the operating system page cache for pooled file reads. + /// + /// This option is available only on Linux. Unaligned logical reads are widened to satisfy the + /// filesystem's direct-I/O requirements and sliced back to the requested range after transfer + /// to the device. + #[cfg(target_os = "linux")] + pub fn with_direct_io(mut self) -> Self { + self.direct_io = true; + self + } +} + +struct PooledHostRead { + buffer: PooledPinnedBuffer, + requested_range: Range, +} + +trait FileReadBackend: Send + Sync { + fn size(&self) -> VortexResult; + + fn read( + &self, + pool: &Arc, + offset: u64, + length: usize, + ) -> VortexResult; +} + +struct BufferedFileReadBackend { + file: File, +} + +impl BufferedFileReadBackend { + fn open(path: &Path) -> VortexResult { + Ok(Self { + file: File::open(path)?, + }) + } +} + +impl FileReadBackend for BufferedFileReadBackend { + fn size(&self) -> VortexResult { + Ok(self.file.metadata()?.len()) + } + + fn read( + &self, + pool: &Arc, + offset: u64, + length: usize, + ) -> VortexResult { + let mut buffer = pool.get(length)?; + read_exact_at(&self.file, buffer.as_mut_slice(), offset)?; + Ok(PooledHostRead { + buffer, + requested_range: 0..length, + }) + } +} + +#[cfg(target_os = "linux")] +fn open_backend( + path: &Path, + options: PooledFileReadAtOptions, +) -> VortexResult> { + if options.direct_io { + Ok(Arc::new(DirectFileReadBackend::open(path)?)) + } else { + Ok(Arc::new(BufferedFileReadBackend::open(path)?)) + } +} + +#[cfg(not(target_os = "linux"))] +fn open_backend( + path: &Path, + _options: PooledFileReadAtOptions, +) -> VortexResult> { + Ok(Arc::new(BufferedFileReadBackend::open(path)?)) +} + +/// File reader that uses CUDA pinned host memory for I/O buffers and transfers +/// directly to the GPU. +/// +/// Reads into a pooled pinned (page-locked) buffer, then submits a non-blocking +/// H2D DMA transfer and returns a device `BufferHandle`. +/// +/// This is a data-plane reader. To open a complete local Vortex file, prefer +/// [`crate::CudaOpenOptionsExt::with_cuda`], which keeps the footer and zone maps on the host. +#[derive(Clone)] +pub struct PooledFileReadAt { + uri: Arc, + backend: Arc, + handle: Handle, + pool: Arc, + stream: VortexCudaStream, +} + +impl PooledFileReadAt { + /// Open a file for pooled reading with direct device transfer. + pub fn open( + path: impl AsRef, + handle: Handle, + pool: Arc, + stream: VortexCudaStream, + ) -> VortexResult { + Self::open_with_options( + path, + handle, + pool, + stream, + PooledFileReadAtOptions::default(), + ) + } + + /// Open a file for pooled reading with explicit options. + pub fn open_with_options( + path: impl AsRef, + handle: Handle, + pool: Arc, + stream: VortexCudaStream, + options: PooledFileReadAtOptions, + ) -> VortexResult { + let path = path.as_ref(); + let uri = Arc::from(path.to_string_lossy().to_string()); + let backend = open_backend(path, options)?; + Ok(Self { + uri, + backend, + handle, + pool, + stream, + }) + } +} + +impl VortexReadAt for PooledFileReadAt { + fn uri(&self) -> Option<&Arc> { + Some(&self.uri) + } + + fn coalesce_config(&self) -> Option { + Some(CoalesceConfig::file()) + } + + fn concurrency(&self) -> usize { + DEFAULT_FILE_CONCURRENCY + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + let backend = Arc::clone(&self.backend); + async move { backend.size() }.boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + let backend = Arc::clone(&self.backend); + let handle = self.handle.clone(); + let stream = self.stream.clone(); + let pool = Arc::clone(&self.pool); + + async move { + let read = handle + .spawn_blocking(move || backend.read(&pool, offset, length)) + .await?; + let cuda_buf = read.buffer.transfer_to_device(&stream)?; + Ok(BufferHandle::new_device(Arc::new(cuda_buf)).slice(read.requested_range)) + } + .boxed() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn pooled_file_read_options_default_to_buffered_io() { + assert!(!PooledFileReadAtOptions::default().direct_io); + } + + #[cfg(target_os = "linux")] + #[test] + fn pooled_file_read_options_enable_direct_io() { + assert!( + PooledFileReadAtOptions::default() + .with_direct_io() + .direct_io + ); + } +} diff --git a/vortex-cuda/src/pooled_read_at/file/direct.rs b/vortex-cuda/src/pooled_read_at/file/direct.rs new file mode 100644 index 00000000000..53366c36915 --- /dev/null +++ b/vortex-cuda/src/pooled_read_at/file/direct.rs @@ -0,0 +1,288 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fs::File; +use std::io; +use std::ops::Range; +use std::os::unix::fs::FileExt; +use std::path::Path; +use std::sync::Arc; + +use rustix::fs::AtFlags; +use rustix::fs::Mode; +use rustix::fs::OFlags; +use rustix::fs::StatxFlags; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; + +use super::FileReadBackend; +use super::PooledHostRead; +use crate::pinned::PinnedByteBufferPool; + +/// Conservative direct-I/O alignment used when Linux cannot report the filesystem constraints. +/// +/// A page-sized fallback is accepted by common block devices and filesystems. If the actual +/// requirement is stricter, the read fails with the underlying `EINVAL`. +const FALLBACK_DIRECT_IO_ALIGNMENT: usize = 4096; + +#[derive(Clone, Copy)] +struct DirectIoConstraints { + /// Required alignment of the address of the userspace I/O buffer. + memory_alignment: usize, + /// Required alignment of both the file offset and the I/O length. + offset_alignment: usize, +} + +impl Default for DirectIoConstraints { + fn default() -> Self { + Self { + memory_alignment: FALLBACK_DIRECT_IO_ALIGNMENT, + offset_alignment: FALLBACK_DIRECT_IO_ALIGNMENT, + } + } +} + +#[derive(Debug, PartialEq, Eq)] +struct DirectIoRange { + read_offset: u64, + read_length: usize, + requested_range: Range, +} + +pub(super) struct DirectFileReadBackend { + file: File, + constraints: DirectIoConstraints, +} + +impl DirectFileReadBackend { + pub(super) fn open(path: &Path) -> VortexResult { + let file = File::from( + rustix::fs::open( + path, + OFlags::RDONLY | OFlags::CLOEXEC | OFlags::DIRECT, + Mode::empty(), + ) + .map_err(io::Error::from)?, + ); + let constraints = direct_io_constraints(&file)?; + Ok(Self { file, constraints }) + } +} + +impl FileReadBackend for DirectFileReadBackend { + fn size(&self) -> VortexResult { + Ok(self.file.metadata()?.len()) + } + + fn read( + &self, + pool: &Arc, + offset: u64, + length: usize, + ) -> VortexResult { + let direct_range = direct_io_range(offset, length, self.constraints.offset_alignment)?; + let mut buffer = pool.get(direct_range.read_length)?; + let address = buffer.as_mut_slice().as_ptr() as usize; + vortex_ensure!( + address.is_multiple_of(self.constraints.memory_alignment), + "pinned buffer address {address:#x} is not aligned to {} bytes", + self.constraints.memory_alignment + ); + + let bytes_read = read_direct_at( + &self.file, + buffer.as_mut_slice(), + direct_range.read_offset, + direct_range.requested_range.end, + self.constraints, + )?; + buffer.truncate(bytes_read); + Ok(PooledHostRead { + buffer, + requested_range: direct_range.requested_range, + }) + } +} + +fn direct_io_range(offset: u64, length: usize, alignment: usize) -> VortexResult { + vortex_ensure!(alignment > 0, "direct I/O alignment must be non-zero"); + if length == 0 { + return Ok(DirectIoRange { + read_offset: offset, + read_length: 0, + requested_range: 0..0, + }); + } + + let alignment_u64 = u64::try_from(alignment)?; + let length_u64 = u64::try_from(length)?; + let requested_end = offset.checked_add(length_u64).ok_or_else(|| { + vortex_err!("direct I/O range overflow: offset={offset}, length={length}") + })?; + let read_offset = offset - offset % alignment_u64; + let read_end = requested_end + .checked_next_multiple_of(alignment_u64) + .ok_or_else(|| vortex_err!("direct I/O aligned end overflow"))?; + let read_length = usize::try_from(read_end - read_offset)?; + let slice_start = usize::try_from(offset - read_offset)?; + let slice_end = slice_start.checked_add(length).ok_or_else(|| { + vortex_err!("direct I/O range overflow: offset={offset}, length={length}") + })?; + + Ok(DirectIoRange { + read_offset, + read_length, + requested_range: slice_start..slice_end, + }) +} + +fn read_direct_at( + file: &File, + buffer: &mut [u8], + offset: u64, + required_bytes: usize, + constraints: DirectIoConstraints, +) -> io::Result { + let mut initialized = 0; + while initialized < required_bytes { + let initialized_u64 = u64::try_from(initialized) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow"))?; + let read_offset = offset + .checked_add(initialized_u64) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow"))?; + let bytes_read = match file.read_at(&mut buffer[initialized..], read_offset) { + Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, + result => result?, + }; + if bytes_read == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "direct read returned {initialized} bytes, but {required_bytes} bytes were required" + ), + )); + } + initialized += bytes_read; + if initialized < required_bytes + && (!initialized.is_multiple_of(constraints.offset_alignment) + || !initialized.is_multiple_of(constraints.memory_alignment)) + { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "direct read returned an unaligned short read of {initialized} bytes before the required {required_bytes} bytes" + ), + )); + } + } + + Ok(initialized) +} + +fn direct_io_constraints(file: &File) -> VortexResult { + let Ok(stat) = rustix::fs::statx( + file, + c"", + AtFlags::EMPTY_PATH | AtFlags::STATX_DONT_SYNC, + StatxFlags::DIOALIGN, + ) else { + return Ok(DirectIoConstraints::default()); + }; + if stat.stx_mask & StatxFlags::DIOALIGN.bits() == 0 { + return Ok(DirectIoConstraints::default()); + } + + let Ok(memory_alignment) = usize::try_from(stat.stx_dio_mem_align) else { + return Ok(DirectIoConstraints::default()); + }; + let Ok(offset_alignment) = usize::try_from(stat.stx_dio_offset_align) else { + return Ok(DirectIoConstraints::default()); + }; + if memory_alignment == 0 || offset_alignment == 0 { + return Ok(DirectIoConstraints::default()); + } + vortex_ensure!( + memory_alignment.is_power_of_two(), + "direct I/O memory alignment must be a power of two, got {memory_alignment}" + ); + vortex_ensure!( + offset_alignment.is_power_of_two(), + "direct I/O offset alignment must be a power of two, got {offset_alignment}" + ); + + Ok(DirectIoConstraints { + memory_alignment, + offset_alignment, + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case(0, 0, 4096, 0, 0, 0)] + #[case(5, 0, 4096, 5, 0, 0)] + #[case(5, 10, 4096, 0, 4096, 5)] + #[case(4090, 20, 4096, 0, 8192, 4090)] + #[case(4096, 4096, 4096, 4096, 4096, 0)] + #[case(513, 1, 512, 512, 512, 1)] + #[case(4096, 8193, 4096, 4096, 12288, 0)] + fn widens_direct_read_to_block_boundaries( + #[case] offset: u64, + #[case] length: usize, + #[case] alignment: usize, + #[case] expected_offset: u64, + #[case] expected_length: usize, + #[case] expected_prefix: usize, + ) -> VortexResult<()> { + assert_eq!( + direct_io_range(offset, length, alignment)?, + DirectIoRange { + read_offset: expected_offset, + read_length: expected_length, + requested_range: expected_prefix..expected_prefix + length, + } + ); + Ok(()) + } + + #[rstest] + #[case(u64::MAX, 2, 4096)] + #[case(0, 1, 0)] + fn rejects_invalid_direct_read_range( + #[case] offset: u64, + #[case] length: usize, + #[case] alignment: usize, + ) { + assert!(direct_io_range(offset, length, alignment).is_err()); + } + + #[test] + fn aligned_ranges_cover_requested_bytes() -> VortexResult<()> { + for alignment in [512, 4096] { + for offset in 0..alignment * 2 { + for length in [0, 1, alignment - 1, alignment, alignment + 1] { + let range = direct_io_range(offset as u64, length, alignment)?; + if length == 0 { + assert_eq!(range.read_length, 0); + continue; + } + + assert_eq!(range.read_offset % alignment as u64, 0); + assert_eq!(range.read_length % alignment, 0); + assert_eq!(range.requested_range.len(), length); + assert!(range.requested_range.end <= range.read_length); + assert_eq!( + range.read_offset + range.requested_range.start as u64, + offset as u64 + ); + } + } + } + Ok(()) + } +}