diff --git a/Cargo.lock b/Cargo.lock index 9a7274a774b..fee1e3ebd34 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10235,6 +10235,7 @@ dependencies = [ "async-stream", "async-trait", "bytes", + "crossbeam-channel", "custom-labels", "futures", "glob", @@ -10464,7 +10465,6 @@ dependencies = [ "arrow-array 58.4.0", "arrow-data 58.4.0", "arrow-schema 58.4.0", - "async-fs", "bytes", "itertools 0.14.0", "log", diff --git a/Cargo.toml b/Cargo.toml index 3ae5f7bdbce..65452f18300 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -129,6 +129,7 @@ cfg-if = "1.0.1" chrono = "0.4.44" clap = "4.5" criterion = { package = "codspeed-criterion-compat-walltime", version = "5.0.0" } +crossbeam-channel = "0.5.16" crossterm = "0.29" cudarc = { version = "0.19.0", features = [ # The NSight Compute version available on lambda.ai hosts does not inject a symbol for diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index b3f6448484a..d6ed23f96d0 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -21,6 +21,7 @@ async-fs = { workspace = true } async-stream = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } +crossbeam-channel = { workspace = true } futures = { workspace = true, features = ["std", "executor"] } glob = { workspace = true } kanal = { workspace = true } diff --git a/vortex-io/src/runtime/blocking_pool.rs b/vortex-io/src/runtime/blocking_pool.rs new file mode 100644 index 00000000000..fc315844831 --- /dev/null +++ b/vortex-io/src/runtime/blocking_pool.rs @@ -0,0 +1,312 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use crossbeam_channel::Receiver; +use crossbeam_channel::RecvTimeoutError; +use crossbeam_channel::Sender; +use crossbeam_channel::TryRecvError; +use crossbeam_channel::unbounded; +use parking_lot::Mutex; +use vortex_error::vortex_panic; + +use crate::runtime::AbortHandle; +use crate::runtime::AbortHandleRef; + +const DEFAULT_MAX_THREADS: usize = 500; +const MIN_MAX_THREADS: usize = 1; +const MAX_MAX_THREADS: usize = 10_000; +const KEEP_ALIVE: Duration = Duration::from_millis(500); + +/// A dynamically sized pool for blocking I/O owned by a single Vortex runtime. +pub(crate) struct BlockingPool { + sender: Sender, + receiver: Receiver, + state: Arc>, + thread_limit: usize, +} + +impl Default for BlockingPool { + fn default() -> Self { + Self::new(max_threads()) + } +} + +impl BlockingPool { + fn new(thread_limit: usize) -> Self { + assert!(thread_limit > 0, "blocking thread limit must be non-zero"); + let (sender, receiver) = unbounded(); + Self { + sender, + receiver, + state: Arc::new(Mutex::new(PoolState::default())), + thread_limit, + } + } + + pub(crate) fn spawn(&self, task: Box) -> AbortHandleRef { + let cancelled = Arc::new(AtomicBool::new(false)); + let mut state = self.state.lock(); + if self + .sender + .send(Job { + cancelled: Arc::clone(&cancelled), + task, + }) + .is_err() + { + vortex_panic!("cannot spawn blocking work on a shut down runtime"); + } + state.queued_job_count += 1; + if state.queued_job_count > state.idle_thread_count { + self.grow(&mut state); + } + Box::new(BlockingAbortHandle { cancelled }) + } + + fn grow(&self, state: &mut PoolState) { + if state.thread_count >= self.thread_limit { + return; + } + state.thread_count += 1; + + let receiver = self.receiver.clone(); + let worker_state = Arc::clone(&self.state); + if let Err(error) = std::thread::Builder::new() + .name("vortex-blocking-io".to_string()) + .spawn(move || worker_loop(receiver, worker_state)) + { + state.thread_count -= 1; + vortex_panic!("failed to spawn a blocking I/O worker: {error}"); + } + } +} + +#[derive(Default)] +struct PoolState { + // All three counters are updated while holding the pool mutex. Keeping queued work and idle + // capacity in the same state makes thread growth independent of channel timing. + thread_count: usize, + idle_thread_count: usize, + queued_job_count: usize, +} + +struct Job { + cancelled: Arc, + task: Box, +} + +impl Job { + fn run(self) { + if !self.cancelled.load(Ordering::Acquire) { + (self.task)(); + } + } +} + +struct BlockingAbortHandle { + cancelled: Arc, +} + +impl AbortHandle for BlockingAbortHandle { + fn abort(self: Box) { + self.cancelled.store(true, Ordering::Release); + } +} + +fn worker_loop(receiver: Receiver, state: Arc>) { + loop { + let mut pool_state = state.lock(); + match receiver.try_recv() { + Ok(job) => { + pool_state.queued_job_count -= 1; + drop(pool_state); + run_job(job); + continue; + } + Err(TryRecvError::Disconnected) => { + pool_state.thread_count -= 1; + return; + } + Err(TryRecvError::Empty) => { + pool_state.idle_thread_count += 1; + } + } + drop(pool_state); + + match receiver.recv_timeout(KEEP_ALIVE) { + Ok(job) => { + let mut pool_state = state.lock(); + pool_state.idle_thread_count -= 1; + pool_state.queued_job_count -= 1; + drop(pool_state); + run_job(job); + } + Err(RecvTimeoutError::Timeout) => { + let mut pool_state = state.lock(); + match receiver.try_recv() { + Ok(job) => { + pool_state.idle_thread_count -= 1; + pool_state.queued_job_count -= 1; + drop(pool_state); + run_job(job); + } + Err(TryRecvError::Empty | TryRecvError::Disconnected) => { + pool_state.idle_thread_count -= 1; + pool_state.thread_count -= 1; + return; + } + } + } + Err(RecvTimeoutError::Disconnected) => { + let mut pool_state = state.lock(); + pool_state.idle_thread_count -= 1; + pool_state.thread_count -= 1; + return; + } + } + } +} + +fn run_job(job: Job) { + // `Handle::spawn_blocking` catches task panics so they can be propagated to the caller. Keep + // this boundary too, so an unexpected panic does not corrupt the worker counts and permanently + // reduce the pool's capacity. + drop(catch_unwind(AssertUnwindSafe(|| job.run()))); +} + +fn max_threads() -> usize { + std::env::var("BLOCKING_MAX_THREADS") + .ok() + .and_then(|value| value.parse::().ok()) + .map(|value| value.clamp(MIN_MAX_THREADS, MAX_MAX_THREADS)) + .unwrap_or(DEFAULT_MAX_THREADS) +} + +#[cfg(test)] +mod tests { + use std::sync::Barrier; + use std::sync::mpsc; + use std::time::Instant; + + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::*; + + #[test] + fn test_executes_jobs() -> VortexResult<()> { + let pool = BlockingPool::new(2); + let (send, recv) = mpsc::sync_channel(1); + drop(pool.spawn(Box::new(move || { + let _ = send.send(42); + }))); + + assert_eq!( + recv.recv_timeout(Duration::from_secs(5)) + .map_err(|error| vortex_err!("blocking job did not finish: {error}"))?, + 42 + ); + Ok(()) + } + + #[test] + fn test_aborts_queued_jobs() -> VortexResult<()> { + let pool = BlockingPool::new(1); + let (started_send, started_recv) = mpsc::sync_channel(1); + let (release_send, release_recv) = mpsc::sync_channel(1); + drop(pool.spawn(Box::new(move || { + let _ = started_send.send(()); + let _ = release_recv.recv(); + }))); + started_recv + .recv_timeout(Duration::from_secs(5)) + .map_err(|error| vortex_err!("first blocking job did not start: {error}"))?; + + let (ran_send, ran_recv) = mpsc::sync_channel(1); + pool.spawn(Box::new(move || { + let _ = ran_send.send(()); + })) + .abort(); + + let (done_send, done_recv) = mpsc::sync_channel(1); + drop(pool.spawn(Box::new(move || { + let _ = done_send.send(()); + }))); + let _ = release_send.send(()); + done_recv + .recv_timeout(Duration::from_secs(5)) + .map_err(|error| vortex_err!("blocking queue did not drain: {error}"))?; + assert!(ran_recv.try_recv().is_err()); + Ok(()) + } + + #[test] + fn test_grows_for_concurrent_jobs() -> VortexResult<()> { + let pool = BlockingPool::new(2); + let barrier = Arc::new(Barrier::new(3)); + let (started_send, started_recv) = mpsc::sync_channel(2); + + for _ in 0..2 { + let barrier = Arc::clone(&barrier); + let started_send = started_send.clone(); + drop(pool.spawn(Box::new(move || { + let _ = started_send.send(()); + barrier.wait(); + }))); + } + + for _ in 0..2 { + started_recv + .recv_timeout(Duration::from_secs(5)) + .map_err(|error| vortex_err!("blocking job did not start: {error}"))?; + } + barrier.wait(); + Ok(()) + } + + #[test] + fn test_reuses_idle_worker() -> VortexResult<()> { + let pool = BlockingPool::new(2); + let (first_send, first_recv) = mpsc::sync_channel(1); + drop(pool.spawn(Box::new(move || { + let _ = first_send.send(std::thread::current().id()); + }))); + let first_thread = first_recv + .recv_timeout(Duration::from_secs(5)) + .map_err(|error| vortex_err!("first blocking job did not finish: {error}"))?; + + let deadline = Instant::now() + Duration::from_secs(5); + while pool.state.lock().idle_thread_count != 1 { + if Instant::now() >= deadline { + return Err(vortex_err!("blocking worker did not become idle")); + } + std::thread::yield_now(); + } + + let (started_send, started_recv) = mpsc::sync_channel(1); + let (release_send, release_recv) = mpsc::sync_channel(1); + drop(pool.spawn(Box::new(move || { + let _ = started_send.send(std::thread::current().id()); + let _ = release_recv.recv(); + }))); + let second_thread = started_recv + .recv_timeout(Duration::from_secs(5)) + .map_err(|error| vortex_err!("second blocking job did not finish: {error}"))?; + + assert_eq!(first_thread, second_thread); + let state = pool.state.lock(); + assert_eq!(state.thread_count, 1); + assert_eq!(state.idle_thread_count, 0); + assert_eq!(state.queued_job_count, 0); + drop(state); + let _ = release_send.send(()); + Ok(()) + } +} diff --git a/vortex-io/src/runtime/current.rs b/vortex-io/src/runtime/current.rs index c4efcde3581..bb1c369af77 100644 --- a/vortex-io/src/runtime/current.rs +++ b/vortex-io/src/runtime/current.rs @@ -15,6 +15,7 @@ use crate::runtime::BlockingRuntime; use crate::runtime::Executor; use crate::runtime::Handle; pub use crate::runtime::pool::CurrentThreadWorkerPool; +use crate::runtime::smol::SmolExecutor; /// A current thread runtime allows callers to much more explicitly drive Vortex futures than with /// a Tokio runtime. @@ -31,7 +32,7 @@ pub use crate::runtime::pool::CurrentThreadWorkerPool; /// with the desired number of worker threads that will drive work on behalf of the runtime. #[derive(Clone, Default)] pub struct CurrentThreadRuntime { - executor: Arc>, + executor: Arc, } impl CurrentThreadRuntime { @@ -72,7 +73,7 @@ impl CurrentThreadRuntime { // ready item without every one of them having to become an executor. let capacity = get_available_parallelism().unwrap_or(1).max(1); let (result_tx, result_rx) = kanal::bounded_async(capacity); - let driver = self.executor.spawn(async move { + let driver = self.executor.async_executor().spawn(async move { futures::pin_mut!(stream); while let Some(item) = stream.next().await { // If all receivers are dropped, we stop driving the stream. @@ -120,7 +121,7 @@ impl BlockingRuntime for CurrentThreadRuntime { /// An iterator that wraps up a stream to drive it using the current thread execution. pub struct CurrentThreadIterator<'a, T> { - executor: Arc>, + executor: Arc, stream: BoxStream<'a, T>, } @@ -134,7 +135,7 @@ impl Iterator for CurrentThreadIterator<'_, T> { /// An iterator that drives a stream from multiple threads. pub struct ThreadSafeIterator { - executor: Arc>, + executor: Arc, results: kanal::AsyncReceiver, /// Handle to the task driving the stream. Once the stream ends, the first consumer to /// observe it joins the task so a panic raised while driving the stream is re-raised rather diff --git a/vortex-io/src/runtime/mod.rs b/vortex-io/src/runtime/mod.rs index 864d3c8b2c6..659f545f430 100644 --- a/vortex-io/src/runtime/mod.rs +++ b/vortex-io/src/runtime/mod.rs @@ -16,6 +16,8 @@ use futures::future::BoxFuture; mod blocking; pub use blocking::*; +#[cfg(not(target_arch = "wasm32"))] +mod blocking_pool; mod handle; pub use handle::*; diff --git a/vortex-io/src/runtime/pool.rs b/vortex-io/src/runtime/pool.rs index 71ffb54cae6..a952c7ca435 100644 --- a/vortex-io/src/runtime/pool.rs +++ b/vortex-io/src/runtime/pool.rs @@ -11,14 +11,16 @@ use smol::block_on; use vortex_error::VortexExpect; use vortex_utils::parallelism::get_available_parallelism; +use crate::runtime::smol::SmolExecutor; + #[derive(Clone)] pub struct CurrentThreadWorkerPool { - executor: Arc>, + executor: Arc, state: Arc>, } impl CurrentThreadWorkerPool { - pub(super) fn new(executor: Arc>) -> Self { + pub(super) fn new(executor: Arc) -> Self { Self { executor, state: Arc::new(Mutex::new(PoolState::default())), diff --git a/vortex-io/src/runtime/smol.rs b/vortex-io/src/runtime/smol.rs index 32aba238146..917e858afcf 100644 --- a/vortex-io/src/runtime/smol.rs +++ b/vortex-io/src/runtime/smol.rs @@ -1,26 +1,56 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Deref; + use futures::future::BoxFuture; use crate::runtime::AbortHandle; use crate::runtime::AbortHandleRef; use crate::runtime::Executor; +use crate::runtime::blocking_pool::BlockingPool; + +/// The async executor and instance-owned blocking pool backing a current-thread runtime. +pub(crate) struct SmolExecutor { + executor: smol::Executor<'static>, + blocking_pool: BlockingPool, +} + +impl Default for SmolExecutor { + fn default() -> Self { + Self { + executor: smol::Executor::new(), + blocking_pool: BlockingPool::default(), + } + } +} + +impl SmolExecutor { + pub(crate) fn async_executor(&self) -> &smol::Executor<'static> { + &self.executor + } +} + +impl Deref for SmolExecutor { + type Target = smol::Executor<'static>; + + fn deref(&self) -> &Self::Target { + &self.executor + } +} -// NOTE(ngates): we implement this for a Weak reference to adhere to the constraint that this -// trait should not hold strong references to the underlying runtime. -impl Executor for smol::Executor<'static> { +impl Executor for SmolExecutor { fn spawn(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef { - SmolAbortHandle::new_handle(smol::Executor::spawn(self, fut)) + SmolAbortHandle::new_handle(self.executor.spawn(fut)) } fn spawn_cpu(&self, task: Box) -> AbortHandleRef { // For now, we spawn CPU work back onto the same execution. - SmolAbortHandle::new_handle(smol::Executor::spawn(self, async move { task() })) + SmolAbortHandle::new_handle(self.executor.spawn(async move { task() })) } fn spawn_blocking_io(&self, task: Box) -> AbortHandleRef { - SmolAbortHandle::new_handle(smol::unblock(task)) + self.blocking_pool.spawn(task) } } diff --git a/vortex-io/src/std_file/mod.rs b/vortex-io/src/std_file/mod.rs index c2e4b12bf40..c30248ed496 100644 --- a/vortex-io/src/std_file/mod.rs +++ b/vortex-io/src/std_file/mod.rs @@ -2,5 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod read_at; +mod write; pub use read_at::*; +pub use write::*; diff --git a/vortex-io/src/std_file/write.rs b/vortex-io/src/std_file/write.rs new file mode 100644 index 00000000000..b4c327cf332 --- /dev/null +++ b/vortex-io/src/std_file/write.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fs::File; +use std::future::Future; +use std::io; +use std::io::Write; +use std::path::Path; +use std::sync::Arc; + +use parking_lot::Mutex; + +use crate::IoBuf; +use crate::VortexWrite; +use crate::runtime::Handle; + +/// A local file writer that runs blocking filesystem operations on a Vortex runtime. +pub struct FileWrite { + file: Arc>, + handle: Handle, +} + +impl FileWrite { + /// Create or truncate a file using `handle` for the blocking open operation. + pub async fn create(path: impl AsRef, handle: Handle) -> io::Result { + let path = path.as_ref().to_path_buf(); + let file = handle.spawn_blocking(move || File::create(path)).await?; + Ok(Self { + file: Arc::new(Mutex::new(file)), + handle, + }) + } +} + +impl VortexWrite for FileWrite { + fn write_all(&mut self, buffer: B) -> impl Future> { + let file = Arc::clone(&self.file); + let handle = self.handle.clone(); + async move { + handle + .spawn_blocking(move || { + file.lock().write_all(buffer.as_slice())?; + Ok(buffer) + }) + .await + } + } + + fn flush(&mut self) -> impl Future> { + let file = Arc::clone(&self.file); + let handle = self.handle.clone(); + async move { handle.spawn_blocking(move || file.lock().flush()).await } + } + + fn shutdown(&mut self) -> impl Future> { + self.flush() + } +} + +#[cfg(test)] +mod tests { + use tempfile::tempdir; + use vortex_error::VortexResult; + + use super::*; + use crate::runtime::BlockingRuntime; + use crate::runtime::current::CurrentThreadRuntime; + + #[test] + fn test_write_file() -> VortexResult<()> { + let directory = tempdir()?; + let path = directory.path().join("write.vortex"); + let runtime = CurrentThreadRuntime::new(); + runtime.block_on(async { + let mut writer = FileWrite::create(&path, runtime.handle()).await?; + writer.write_all(vec![1, 2, 3]).await?; + writer.shutdown().await + })?; + + assert_eq!(std::fs::read(path)?, vec![1, 2, 3]); + Ok(()) + } +} diff --git a/vortex-python/Cargo.toml b/vortex-python/Cargo.toml index b7a3f6e3588..84fe2129203 100644 --- a/vortex-python/Cargo.toml +++ b/vortex-python/Cargo.toml @@ -33,7 +33,6 @@ tui = ["dep:tokio", "dep:vortex-tui"] arrow-array = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true } -async-fs = { workspace = true } bytes = { workspace = true } itertools = { workspace = true } log = { workspace = true } diff --git a/vortex-python/python/vortex/_lib/file.pyi b/vortex-python/python/vortex/_lib/file.pyi index f8c928be459..901e0323eec 100644 --- a/vortex-python/python/vortex/_lib/file.pyi +++ b/vortex-python/python/vortex/_lib/file.pyi @@ -5,6 +5,7 @@ from typing import final import polars as pl import pyarrow as pa +from typing_extensions import override from vortex.type_aliases import IntoProjection @@ -22,6 +23,10 @@ class VortexFile: def __len__(self) -> int: ... @property def dtype(self) -> DType: ... + @property + def path(self) -> str: ... + @override + def __reduce__(self) -> tuple[object, tuple[str, bool]]: ... def scan( self, projection: IntoProjection = None, diff --git a/vortex-python/python/vortex/datasets.py b/vortex-python/python/vortex/datasets.py index 3e5536fde1e..a09da332def 100644 --- a/vortex-python/python/vortex/datasets.py +++ b/vortex-python/python/vortex/datasets.py @@ -48,7 +48,6 @@ _BaseExamplesIterable, get_format_type_from_alias, ) - from datasets.table import InMemoryTable from huggingface_hub import HfApi, snapshot_download except ImportError as e: # pragma: no cover - exercised only without optional deps. raise ImportError("Install vortex-data[hf] to use vortex.datasets.") from e @@ -528,20 +527,9 @@ def _materialize_dataset( num_proc: int | None, ) -> hf_datasets.Dataset: features = _features_for_files(files, _normalize_columns(columns)) - if filter is not None: - # vortex.Expr cannot be pickled, so it cannot pass through Dataset.from_generator's - # gen_kwargs (which Hugging Face hashes for the cache fingerprint). Read the filtered rows - # in-process and build an in-memory dataset from the resulting Arrow tables instead; - # cache_dir, keep_in_memory, and num_proc do not apply to this path. - return _materialize_filtered( - files, - columns=columns, - filter=filter, - limit=limit, - batch_size=batch_size, - split=split, - features=features, - ) + # `vortex.Expr` is picklable (via its protobuf wire format), so a filter passes through + # `Dataset.from_generator`'s gen_kwargs -- which Hugging Face both hashes for the cache + # fingerprint and ships to `num_proc` worker processes -- like any other argument. gen_kwargs = { "files": list(files), # Keep `columns` a tuple (never a list): Hugging Face shards `num_proc` work across the @@ -567,44 +555,6 @@ def _materialize_dataset( ) -def _materialize_filtered( - files: Sequence[str], - *, - columns: Sequence[str] | None, - filter: Expr | None, - limit: int | None, - batch_size: int | None, - split: str, - features: hf_datasets.Features, -) -> hf_datasets.Dataset: - tables: list[pa.Table] = [] - yielded = 0 - for file_name in files: - if limit is not None and yielded >= limit: - break - for table in _scan_file_as_tables( - file_name, - columns=_normalize_columns(columns), - filter=filter, - limit=None if limit is None else limit - yielded, - batch_size=batch_size, - ): - if limit is not None and yielded + len(table) > limit: - table = table.slice(0, limit - yielded) - if len(table) == 0: - continue - tables.append(table) - yielded += len(table) - if limit is not None and yielded >= limit: - break - combined = pa.concat_tables(tables) if tables else features.arrow_schema.empty_table() - return hf_datasets.Dataset( - InMemoryTable(combined), - info=hf_datasets.DatasetInfo(features=features), - split=hf_datasets.Split(split), - ) - - def _generate_rows( files: Sequence[str], columns: Sequence[str] | None, diff --git a/vortex-python/python/vortex/file.py b/vortex-python/python/vortex/file.py index 7ca666acf50..4db19c1c05b 100644 --- a/vortex-python/python/vortex/file.py +++ b/vortex-python/python/vortex/file.py @@ -77,6 +77,11 @@ def dtype(self) -> DType: """The dtype of the file.""" return self._file.dtype + @property + def path(self) -> str: + """The path or URL this file was opened from.""" + return self._file.path + def splits(self) -> list[tuple[int, int]]: return self._file.splits() diff --git a/vortex-python/src/dataset.rs b/vortex-python/src/dataset.rs index 2177e56d137..9b6aaea7132 100644 --- a/vortex-python/src/dataset.rs +++ b/vortex-python/src/dataset.rs @@ -28,10 +28,10 @@ use vortex::layout::scan::split_by::SplitBy; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_arrow::ArrowSessionExt; -use crate::RUNTIME; use crate::arrays::PyArrayRef; use crate::arrow::IntoPyArrow; use crate::arrow::ToPyArrow; +use crate::current_runtime; use crate::error::PyVortexResult; use crate::expr::PyExpr; use crate::install_module; @@ -85,7 +85,8 @@ pub fn read_array_from_reader( scan = scan.with_row_range(l..r); } - scan.into_array_iter(&*RUNTIME)?.read_all() + let runtime = current_runtime(); + scan.into_array_iter(&runtime)?.read_all() } fn projection_from_python(columns: Option>>) -> PyResult { @@ -212,8 +213,9 @@ impl PyVortexDataset { } let schema = Arc::new(session().arrow().to_arrow_schema(&scan.dtype()?)?); + let runtime = current_runtime(); let reader: Box = - Box::new(scan.into_record_batch_reader(schema, &*RUNTIME)?); + Box::new(scan.into_record_batch_reader(schema, &runtime)?); VortexResult::Ok(reader) })?; @@ -256,7 +258,8 @@ impl PyVortexDataset { scan = scan.with_row_range(l..r); } - scan.into_array_iter(&*RUNTIME)? + let runtime = current_runtime(); + scan.into_array_iter(&runtime)? .map_ok(|array| array.len()) .process_results(|iter| iter.sum()) })?; @@ -290,5 +293,5 @@ pub fn dataset_from_url( None }; - Ok(py.detach(move || RUNTIME.block_on(PyVortexDataset::from_url(url, store_arc)))?) + Ok(py.detach(move || current_runtime().block_on(PyVortexDataset::from_url(url, store_arc)))?) } diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index e5cc371827e..72c3fefb877 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use arrow_array::RecordBatchReader; use arrow_schema::Schema; use pyo3::exceptions::PyTypeError; +use pyo3::intern; use pyo3::prelude::*; use pyo3::types::PyList; use vortex::array::ArrayRef; @@ -30,10 +31,10 @@ use vortex::layout::segments::MokaSegmentCache; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_arrow::ArrowSessionExt; -use crate::RUNTIME; use crate::arrays::PyArrayRef; use crate::arrow::FromPyArrow; use crate::arrow::IntoPyArrow; +use crate::current_runtime; use crate::dataset::PyVortexDataset; use crate::dtype::PyDType; use crate::error::PyVortexResult; @@ -52,11 +53,18 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { install_module("vortex._lib.file", &m)?; m.add_function(wrap_pyfunction!(open, &m)?)?; + m.add_function(wrap_pyfunction!(_reopen, &m)?)?; m.add_class::()?; Ok(()) } +/// Reopen a Vortex file by path. The unpickling half of [`PyVortexFile::__reduce__`]. +#[pyfunction] +fn _reopen(py: Python, path: &str, without_segment_cache: bool) -> PyVortexResult { + open(py, path, None, without_segment_cache) +} + /// Open a Vortex file for reading. /// /// Callers can optionally configure an object store to build from using one of the definitions @@ -69,8 +77,10 @@ pub fn open( store: Option, without_segment_cache: bool, ) -> PyVortexResult { + let had_store = store.is_some(); + let owned_path = path.to_string(); let vxf = py.detach(move || { - RUNTIME.block_on(async move { + current_runtime().block_on(async move { let mut options = session().open_options(); if !without_segment_cache { // TODO(ngates): use a globally shared segment cache for all files @@ -86,12 +96,23 @@ pub fn open( }) })?; - Ok(PyVortexFile { vxf }) + Ok(PyVortexFile { + vxf, + path: owned_path, + had_store, + without_segment_cache, + }) } #[pyclass(name = "VortexFile", module = "vortex", frozen)] pub struct PyVortexFile { vxf: VortexFile, + /// The path this file was opened from, retained so that it can be reopened in another process. + path: String, + /// Whether an explicit object store was passed to [`open`]. Object stores are not picklable, so + /// such a file cannot be reopened by path alone. + had_store: bool, + without_segment_cache: bool, } #[pymethods] @@ -100,6 +121,39 @@ impl PyVortexFile { Ok(usize::try_from(slf.vxf.row_count())?) } + /// The path or URL this file was opened from. + /// + /// Returns + /// ------- + /// :class:`.str` + #[getter] + fn path(slf: PyRef) -> String { + slf.path.clone() + } + + /// Support for Python's pickle protocol: the file is reopened by path in the receiving process. + /// + /// Only the path is transferred, never any read state or segment cache, so this is cheap enough + /// to send a file to every worker of a multiprocessing pool or Ray job. + /// + /// Raises + /// ------ + /// :class:`TypeError` + /// If the file was opened with an explicit ``store``, since object stores cannot be + /// pickled. Pass the URL to :func:`vortex.open` in the worker instead. + fn __reduce__<'py>(slf: PyRef<'py, Self>) -> PyResult<(Bound<'py, PyAny>, (String, bool))> { + if slf.had_store { + return Err(PyTypeError::new_err( + "cannot pickle a VortexFile opened with an explicit store, because object stores \ + are not picklable; open it from its URL in the receiving process instead", + )); + } + let py = slf.py(); + let module = PyModule::import(py, "vortex._lib.file")?; + let reopen = module.getattr(intern!(py, "_reopen"))?; + Ok((reopen, (slf.path.clone(), slf.without_segment_cache))) + } + #[getter] fn dtype(slf: Bound) -> PyResult> { PyDType::init(slf.py(), slf.get().vxf.dtype().clone()) @@ -124,8 +178,9 @@ impl PyVortexFile { let mut ctx = session.create_execution_ctx(); let builder = scan_builder(&vxf, projection, expr, limit, indices, batch_size, &mut ctx)?; + let runtime = current_runtime(); Ok(PyArrayIterator::new(Box::new( - builder.into_array_iter(&*RUNTIME)?, + builder.into_array_iter(&runtime)?, ))) }) } @@ -171,6 +226,7 @@ impl PyVortexFile { .transpose()? .map(Arc::new); + let runtime = current_runtime(); let reader = slf.py().detach(|| { let filter = expr .map(|e| { @@ -201,7 +257,7 @@ impl PyVortexFile { Some(schema) => schema, None => Arc::new(session().arrow().to_arrow_schema(&builder.dtype()?)?), }; - builder.into_record_batch_reader(schema, &*RUNTIME) + builder.into_record_batch_reader(schema, &runtime) })?; let rbr: Box = Box::new(reader); diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 901ab1287c8..2be1025b811 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -5,7 +5,6 @@ use std::sync::Arc; use arrow_array::RecordBatchReader; use arrow_array::ffi_stream::ArrowArrayStreamReader; -use async_fs::File; use pyo3::exceptions::PyTypeError; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; @@ -25,15 +24,16 @@ use vortex::file::WriteStrategyBuilder; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::io::runtime::BlockingRuntime; +use vortex::io::std_file::FileWrite; use vortex_arrow::ArrowSessionExt; use crate::PyVortex; -use crate::RUNTIME; use crate::arrays::PyArray; use crate::arrays::PyArrayRef; use crate::arrow::FromPyArrow; use crate::classes::record_batch_reader_class; use crate::classes::table_class; +use crate::current_runtime; use crate::dataset::PyVortexDataset; use crate::error::PyVortexResult; use crate::expr::PyExpr; @@ -137,7 +137,8 @@ pub fn read_url<'py>( None }; - let dataset = py.detach(move || RUNTIME.block_on(PyVortexDataset::from_url(url, store_arc)))?; + let dataset = + py.detach(move || current_runtime().block_on(PyVortexDataset::from_url(url, store_arc)))?; dataset.to_array_inner(py, projection, row_filter, indices, row_range) } @@ -248,7 +249,7 @@ pub fn write( ) -> PyVortexResult<()> { let session = session(); py.detach(|| { - RUNTIME.block_on(async move { + current_runtime().block_on(async move { match resolve_store(path, store.map(|x| x.into_inner()))? { ResolvedStore::ObjectStore(store, path) => { let mut store = ObjectStoreWrite::new(store, &path).await?; @@ -260,7 +261,7 @@ pub fn write( VortexResult::Ok(()) } ResolvedStore::Path(path) => { - let mut w = File::create(path).await?; + let mut w = FileWrite::create(path, current_runtime().handle()).await?; session .write_options() .write(&mut w, iter.into_inner().into_array_stream()) @@ -383,7 +384,7 @@ impl PyVortexWriteOptions { .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()); } let strategy = strategy.build(); - RUNTIME.block_on(async move { + current_runtime().block_on(async move { match resolve_store(path, store.map(|x| x.into_inner()))? { ResolvedStore::ObjectStore(store, path) => { let mut store = ObjectStoreWrite::new(store, &path).await?; @@ -396,7 +397,7 @@ impl PyVortexWriteOptions { VortexResult::Ok(()) } ResolvedStore::Path(path) => { - let mut w = File::create(path).await?; + let mut w = FileWrite::create(path, current_runtime().handle()).await?; session .write_options() .with_strategy(strategy) diff --git a/vortex-python/src/lib.rs b/vortex-python/src/lib.rs index d171afc1eb4..0183bd772d8 100644 --- a/vortex-python/src/lib.rs +++ b/vortex-python/src/lib.rs @@ -2,7 +2,16 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::ops::Deref; +#[cfg(unix)] +use std::ptr; +#[cfg(not(unix))] use std::sync::LazyLock; +#[cfg(unix)] +use std::sync::OnceLock; +#[cfg(unix)] +use std::sync::atomic::AtomicPtr; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use log::LevelFilter; use pyo3::exceptions::PyRuntimeError; @@ -36,28 +45,137 @@ mod serde; mod session; mod store; +use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; use vortex::io::runtime::current::CurrentThreadWorkerPool; +use vortex::utils::parallelism::get_available_parallelism; -/// Shared current-thread runtime backing Python Vortex operations. -pub(crate) static RUNTIME: LazyLock = - LazyLock::new(CurrentThreadRuntime::new); +#[cfg(unix)] +use crate::session::reset_session_handle; -/// Shared worker pool that drives [`RUNTIME`]'s executor in the background. +/// The current-thread runtime and its background worker pool. +struct RuntimeState { + runtime: CurrentThreadRuntime, + pool: CurrentThreadWorkerPool, +} + +impl RuntimeState { + fn new() -> Self { + let runtime = CurrentThreadRuntime::new(); + let pool = runtime.new_pool(); + pool.set_workers(requested_workers()); + Self { runtime, pool } + } +} + +#[cfg(not(unix))] +static RUNTIME_STATE: LazyLock = LazyLock::new(RuntimeState::new); + +/// A process-tagged initialization cell. /// -/// On first access, the pool is sized to `VORTEX_MAX_THREADS` (if set to a -/// non-negative integer) or otherwise to `available_parallelism() - 1`. -pub(crate) static POOL: LazyLock = LazyLock::new(|| { - let pool = RUNTIME.new_pool(); - match std::env::var("VORTEX_MAX_THREADS") +/// The pid is checked before touching `state`, so a child never waits on a `OnceLock` whose +/// initializer was running on another thread at the instant of `fork(2)`. +#[cfg(unix)] +struct RuntimeSlot { + pid: u32, + state: OnceLock, +} + +#[cfg(unix)] +static RUNTIME_SLOT: AtomicPtr = AtomicPtr::new(ptr::null_mut()); + +/// The worker count to give a newly built pool, or [`WORKERS_UNSET`] to derive one. +/// +/// Held outside the process-local runtime state so that a forked child inherits the parent's +/// configuration even though it discards the parent's runtime. +static REQUESTED_WORKERS: AtomicUsize = AtomicUsize::new(WORKERS_UNSET); + +const WORKERS_UNSET: usize = usize::MAX; + +/// The worker count for a new pool: whatever [`runtime::set_worker_threads`] last requested, +/// else `VORTEX_MAX_THREADS` if it is set to a non-negative integer, else +/// `available_parallelism() - 1`. +fn requested_workers() -> usize { + match REQUESTED_WORKERS.load(Ordering::Relaxed) { + WORKERS_UNSET => {} + workers => return workers, + } + if let Some(n) = std::env::var("VORTEX_MAX_THREADS") .ok() .and_then(|s| s.parse::().ok()) { - Some(n) => pool.set_workers(n), - None => pool.set_workers_to_available_parallelism(), + return n; } - pool -}); + get_available_parallelism() + .map(|n| n.saturating_sub(1).max(1)) + .unwrap_or(1) +} + +/// Record the worker count requested by the user, so that a forked child inherits it. +pub(crate) fn set_requested_workers(workers: usize) { + REQUESTED_WORKERS.store(workers, Ordering::Relaxed); +} + +#[cfg(unix)] +fn runtime_slot() -> &'static RuntimeSlot { + let pid = std::process::id(); + loop { + let current = RUNTIME_SLOT.load(Ordering::Acquire); + if !current.is_null() { + // SAFETY: Published slots are never freed or replaced in-place. A forked child may + // publish a new slot, but deliberately leaks the inherited allocation. + let slot = unsafe { &*current }; + if slot.pid == pid { + return slot; + } + } + + let fresh = Box::into_raw(Box::new(RuntimeSlot { + pid, + state: OnceLock::new(), + })); + match RUNTIME_SLOT.compare_exchange(current, fresh, Ordering::AcqRel, Ordering::Acquire) { + Ok(_) => { + // SAFETY: `fresh` was just published and published slots are never freed. + return unsafe { &*fresh }; + } + Err(_) => { + // SAFETY: The failed compare-exchange proves `fresh` was never published. + drop(unsafe { Box::from_raw(fresh) }); + } + } + } +} + +#[cfg(unix)] +fn runtime_state() -> &'static RuntimeState { + runtime_slot().state.get_or_init(|| { + let state = RuntimeState::new(); + // Existing objects hold clones of the shared session, so repoint it before publishing the + // new runtime state to callers in this process. + reset_session_handle(state.runtime.handle()); + state + }) +} + +#[cfg(not(unix))] +fn runtime_state() -> &'static RuntimeState { + &RUNTIME_STATE +} + +/// The current-thread runtime backing Python Vortex operations. +pub(crate) fn current_runtime() -> CurrentThreadRuntime { + runtime_state().runtime.clone() +} + +/// Runs `f` with the worker pool that drives [`current_runtime`]'s executor in the background. +/// +/// The pool is deliberately not handed out by value: `CurrentThreadWorkerPool` shares its state +/// between clones but shuts every worker down on `Drop`, so dropping a temporary clone would stop +/// the whole pool. +pub(crate) fn with_pool(f: impl FnOnce(&CurrentThreadWorkerPool) -> T) -> T { + f(&runtime_state().pool) +} /// Vortex is an Apache Arrow-compatible toolkit for working with compressed array data. #[cfg(feature = "extension-module")] diff --git a/vortex-python/src/runtime.rs b/vortex-python/src/runtime.rs index 28b54496218..8cee8985412 100644 --- a/vortex-python/src/runtime.rs +++ b/vortex-python/src/runtime.rs @@ -4,8 +4,9 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use crate::POOL; use crate::install_module; +use crate::set_requested_workers; +use crate::with_pool; pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { let m = PyModule::new(py, "runtime")?; @@ -24,22 +25,27 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { #[pyfunction] #[pyo3(signature = (n=None))] pub fn set_worker_threads(n: Option) -> PyResult<()> { - match n { - Some(n) => { - if n < 0 { - return Err(PyValueError::new_err( - "worker thread count must be non-negative", - )); - } - POOL.set_workers(n as usize); - } - None => POOL.set_workers_to_available_parallelism(), + if let Some(n) = n + && n < 0 + { + return Err(PyValueError::new_err( + "worker thread count must be non-negative", + )); } + let workers = with_pool(|pool| { + match n { + Some(n) => pool.set_workers(n as usize), + None => pool.set_workers_to_available_parallelism(), + } + pool.worker_count() + }); + // Remember the request so that a forked child rebuilds its pool with the same size. + set_requested_workers(workers); Ok(()) } /// Return the current number of background worker threads. #[pyfunction] pub fn worker_threads() -> usize { - POOL.worker_count() + with_pool(|pool| pool.worker_count()) } diff --git a/vortex-python/src/scan.rs b/vortex-python/src/scan.rs index cfdf77a6b8b..913443cc298 100644 --- a/vortex-python/src/scan.rs +++ b/vortex-python/src/scan.rs @@ -11,7 +11,7 @@ use vortex::error::VortexResult; use vortex::layout::scan::repeated_scan::RepeatedScan; use vortex::scalar::Scalar; -use crate::RUNTIME; +use crate::current_runtime; use crate::error::PyVortexResult; use crate::install_module; use crate::iter::PyArrayIterator; @@ -52,8 +52,9 @@ impl PyRepeatedScan { let scan = Arc::clone(&slf.get().scan); slf.py().detach(move || { + let runtime = current_runtime(); Ok(PyArrayIterator::new(Box::new( - scan.execute_array_iter(row_range, &*RUNTIME)?, + scan.execute_array_iter(row_range, &runtime)?, ))) }) } @@ -71,7 +72,8 @@ impl PyRepeatedScan { let scan = Arc::clone(&slf.get().scan); let scalar = slf.py().detach(move || -> VortexResult> { let session = session(); - for batch in scan.execute_array_iter(Some(index..index + 1), &*RUNTIME)? { + let runtime = current_runtime(); + for batch in scan.execute_array_iter(Some(index..index + 1), &runtime)? { let array = batch?; if array.is_empty() { continue; diff --git a/vortex-python/src/session.rs b/vortex-python/src/session.rs index 5598241d2ef..651addb5457 100644 --- a/vortex-python/src/session.rs +++ b/vortex-python/src/session.rs @@ -1,18 +1,87 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#[cfg(unix)] +use std::ptr; +#[cfg(not(unix))] use std::sync::LazyLock; +#[cfg(unix)] +use std::sync::atomic::AtomicPtr; +#[cfg(unix)] +use std::sync::atomic::Ordering; use vortex::VortexSessionDefault; use vortex::io::runtime::BlockingRuntime; +#[cfg(unix)] +use vortex::io::runtime::Handle; use vortex::io::session::RuntimeSessionExt; +#[cfg(unix)] +use vortex::session::SessionExt; use vortex::session::VortexSession; -use crate::RUNTIME; +use crate::current_runtime; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_handle(RUNTIME.handle())); +#[cfg(not(unix))] +static SESSION: LazyLock = LazyLock::new(new_session); +/// The shared session is published without an initialization lock so a forked child cannot inherit +/// it in a permanently initializing state. +#[cfg(unix)] +static SESSION: AtomicPtr = AtomicPtr::new(ptr::null_mut()); + +#[cfg(not(unix))] pub(crate) fn session() -> &'static VortexSession { &SESSION } + +#[cfg(unix)] +pub(crate) fn session() -> &'static VortexSession { + // Ensure a forked child has published its new runtime and repointed an existing session before + // returning that session to a caller. + let runtime = current_runtime(); + loop { + let current = SESSION.load(Ordering::Acquire); + if !current.is_null() { + // SAFETY: A published session is never freed or replaced. + return unsafe { &*current }; + } + + let fresh = Box::into_raw(Box::new( + VortexSession::default().with_handle(runtime.handle()), + )); + match SESSION.compare_exchange(current, fresh, Ordering::AcqRel, Ordering::Acquire) { + Ok(_) => { + // SAFETY: `fresh` was just published and published sessions are never freed. + return unsafe { &*fresh }; + } + Err(_) => { + // SAFETY: The failed compare-exchange proves `fresh` was never published. + drop(unsafe { Box::from_raw(fresh) }); + } + } + } +} + +#[cfg(not(unix))] +fn new_session() -> VortexSession { + VortexSession::default().with_handle(current_runtime().handle()) +} + +/// Point the shared session at `handle`, replacing any previously configured runtime handle. +/// +/// A [`Handle`] is a weak reference to its executor, so after the runtime is rebuilt in a forked +/// child (see [`crate::current_runtime`]) the session must be repointed, or every spawn would panic +/// on a dropped runtime. `VortexSession` has interior mutability, so the change is visible through +/// every existing clone of the session. +/// +/// Does nothing if the session has not been built yet — in that case it will pick up the current +/// runtime's handle when it is first built. +#[cfg(unix)] +pub(crate) fn reset_session_handle(handle: Handle) { + let current = SESSION.load(Ordering::Acquire); + if !current.is_null() { + // SAFETY: A published session is never freed or replaced. + let session = unsafe { &*current }; + session.session().with_handle(handle); + } +} diff --git a/vortex-python/test/child_scenarios.py b/vortex-python/test/child_scenarios.py new file mode 100644 index 00000000000..e6f8a9d677f --- /dev/null +++ b/vortex-python/test/child_scenarios.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""The child-process scenarios driven by `test_fork.py`. + +Run as a script, one scenario per interpreter: + + python child_scenarios.py + +It prints `OK` on success and raises otherwise. + +This is a real module rather than text handed to `python -c` because `spawn` and `forkserver` +re-import `__main__` in the child, and can only pickle callables that live at module level. Every +child target below is therefore top-level and takes its arguments explicitly instead of closing +over parent state. +""" + +from __future__ import annotations + +import multiprocessing +import os +import pickle +import sys +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from multiprocessing.context import BaseContext +from typing import TypeVar, cast + +import pyarrow as pa + +import vortex as vx +import vortex.expr as ve +from vortex.expr import Expr +from vortex.file import VortexFile + +ROWS = 10_000 +CHILD_TIMEOUT_SECONDS = 60 + +FORK = "fork" +SPAWN = "spawn" +FORKSERVER = "forkserver" +ALL_START_METHODS = (FORK, SPAWN, FORKSERVER) + +T = TypeVar("T") + + +def in_child(ctx: BaseContext, target: Callable[..., T], *args: object) -> T: + """Run `target` in one child process and return its result, or raise on hang or error.""" + with ctx.Pool(1) as pool: + result = pool.apply_async(target, args) + try: + return result.get(timeout=CHILD_TIMEOUT_SECONDS) + except multiprocessing.TimeoutError: + pool.terminate() + raise AssertionError(f"child produced no result within {CHILD_TIMEOUT_SECONDS}s") from None + + +# Child targets. Top-level so that every start method can pickle them by reference. + + +def count_rows(path: str) -> int: + return len(vx.open(path).scan().read_all()) + + +def count_file_rows(file: VortexFile) -> int: + return len(file.scan().read_all()) + + +def count_filtered(args: tuple[str, Expr]) -> int: + path, expr = args + return len(vx.open(path).scan(expr=expr).read_all()) + + +def filtered_names(path: str, expr: Expr) -> list[str]: + table = vx.open(path).scan(["name"], expr=expr).read_all().to_arrow_table() + return cast("list[str]", table.column("name").to_pylist()) + + +def write_then_count(values: vx.Array, path: str) -> int: + vx.io.write(values, path) + return len(vx.open(path)) + + +def worker_threads() -> int: + return vx.worker_threads() + + +# Scenarios, each run in the parent of a fresh interpreter. + + +def read_after_parent_read(ctx: BaseContext, path: str) -> None: + """The parent has already built a runtime, so the child inherits a populated (and useless) one + under `fork` and must build its own before it can read anything.""" + assert count_rows(path) == ROWS + assert in_child(ctx, count_rows, path) == ROWS + + +def read_after_saturated_blocking_pool(ctx: BaseContext, path: str) -> None: + """Reaching the blocking-pool limit in the parent must not prevent a fresh pool from starting in + the child. The old process-global `blocking` executor could never recover from this state.""" + os.environ["BLOCKING_MAX_THREADS"] = "1" + assert count_rows(path) == ROWS + assert in_child(ctx, count_rows, path) == ROWS + + +def write_after_parent_write(ctx: BaseContext, path: str) -> None: + """Local writes use the runtime-owned blocking pool too; otherwise async-fs would retain the + same process-global fork hazard after the read path was fixed.""" + values = vx.array(pa.array([{"value": i} for i in range(10)])) + vx.io.write(values, f"{path}.parent-write") + assert in_child(ctx, write_then_count, values, f"{path}.child-write") == 10 + + +def child_has_live_workers(ctx: BaseContext, path: str) -> None: + """The child must end up with real worker threads, not the parent's phantom handles.""" + assert count_rows(path) == ROWS + assert in_child(ctx, worker_threads) >= 1 + + +def expr_filter_in_child(ctx: BaseContext, path: str) -> None: + """An Expr pickled into the child and evaluated there.""" + assert count_rows(path) == ROWS + expr = (ve.column("age") >= ROWS - 2) & ve.like(ve.column("name"), "person-%") + assert in_child(ctx, filtered_names, path, expr) == [f"person-{ROWS - 2}", f"person-{ROWS - 1}"] + + +def child_inherits_worker_count(ctx: BaseContext, path: str) -> None: + """The worker count configured in the parent must carry over to the forked child's fresh pool. + + Fork-only: `spawn` and `forkserver` children start from a new interpreter, so they legitimately + know nothing about the parent's configuration. + """ + vx.set_worker_threads(3) + assert count_rows(path) == ROWS + assert in_child(ctx, worker_threads) == 3 + + +def pool_map_with_pickled_expr(ctx: BaseContext, path: str) -> None: + """A Pool pickles both the callable and its arguments: the shape `datasets` uses for `num_proc`, + and the reason Expr needs `__reduce__`.""" + assert count_rows(path) == ROWS + expr = ve.column("age") < 100 + with ctx.Pool(2) as pool: + result = pool.map_async(count_filtered, [(path, expr), (path, expr)]) + try: + counts = result.get(timeout=CHILD_TIMEOUT_SECONDS) + except multiprocessing.TimeoutError: + pool.terminate() + raise AssertionError(f"Pool did not finish within {CHILD_TIMEOUT_SECONDS}s") from None + assert counts == [100, 100], counts + + +def pickled_file_reads_in_child(ctx: BaseContext, path: str) -> None: + """A pickled VortexFile is reopened by path in the child.""" + vxf = vx.open(path) + assert len(vxf.scan().read_all()) == ROWS + restored = cast(VortexFile, cast(object, pickle.loads(pickle.dumps(vxf)))) + assert restored.path == path + assert in_child(ctx, count_file_rows, restored) == ROWS + + +@dataclass(frozen=True) +class Scenario: + run: Callable[[BaseContext, str], None] + start_methods: Sequence[str] = ALL_START_METHODS + + +SCENARIOS: dict[str, Scenario] = { + "read_after_parent_read": Scenario(read_after_parent_read), + "read_after_saturated_blocking_pool": Scenario(read_after_saturated_blocking_pool), + "write_after_parent_write": Scenario(write_after_parent_write), + "child_has_live_workers": Scenario(child_has_live_workers), + "expr_filter_in_child": Scenario(expr_filter_in_child), + "child_inherits_worker_count": Scenario(child_inherits_worker_count, start_methods=(FORK,)), + "pool_map_with_pickled_expr": Scenario(pool_map_with_pickled_expr), + "pickled_file_reads_in_child": Scenario(pickled_file_reads_in_child), +} + + +def main(argv: Sequence[str]) -> None: + name, start_method, path = argv + scenario = SCENARIOS[name] + if start_method not in scenario.start_methods: + raise AssertionError(f"scenario {name!r} does not apply to the {start_method!r} start method") + scenario.run(multiprocessing.get_context(start_method), path) + print("OK") + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/vortex-python/test/test_fork.py b/vortex-python/test/test_fork.py new file mode 100644 index 00000000000..832ad6493af --- /dev/null +++ b/vortex-python/test/test_fork.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""The Vortex runtime must be rebuilt in forked children, or every read there hangs. + +Every scenario in `child_scenarios.py` runs against each `multiprocessing` start method available +on the platform. `fork` is the hazard being guarded against; `spawn` and `forkserver` start their +children from a fresh interpreter, so they mainly prove that the pickling those methods require +(Expr, Array, VortexFile) still works and that a child built from scratch reads just as well. + +Each combination runs in a **fresh interpreter** that imports nothing but Vortex, and only then +creates children. That keeps the tests measuring Vortex's own behaviour: several third-party +libraries (Ray, DuckDB, CUDA drivers) install their own `atfork` handlers or hold locks that make +`fork` unsafe process-wide, so forking straight from the pytest process would be testing them as +much as us. + +A regression shows up as a timeout, since the failure mode being guarded against is a hang. +""" + +from __future__ import annotations + +import multiprocessing +import subprocess +import sys +from pathlib import Path + +import pyarrow as pa +import pytest + +import vortex as vx + +from .child_scenarios import ROWS, SCENARIOS + +TIMEOUT_SECONDS = 120 +SCRIPT = Path(__file__).with_name("child_scenarios.py") +AVAILABLE_START_METHODS = frozenset(multiprocessing.get_all_start_methods()) + + +SCENARIO_CASES: list[tuple[str, str]] = [ + (name, start_method) for name in sorted(SCENARIOS) for start_method in SCENARIOS[name].start_methods +] + + +@pytest.fixture(scope="module") +def people(tmp_path_factory: pytest.TempPathFactory) -> Path: + path = tmp_path_factory.mktemp("fork") / "people.vortex" + vx.io.write(vx.array(pa.array([{"name": f"person-{i}", "age": i} for i in range(ROWS)])), str(path)) + return path + + +@pytest.mark.parametrize(("name", "start_method"), SCENARIO_CASES) +def test_child_process_scenario(name: str, start_method: str, people: Path) -> None: + if start_method not in AVAILABLE_START_METHODS: + pytest.skip(f"the {start_method!r} start method is not available on {sys.platform}") + + try: + proc = subprocess.run( + [sys.executable, str(SCRIPT), name, start_method, str(people)], + capture_output=True, + text=True, + timeout=TIMEOUT_SECONDS, + check=False, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"scenario {name!r} ({start_method}) did not finish within {TIMEOUT_SECONDS}s (likely a deadlock)") + + if proc.returncode != 0 or not proc.stdout.strip().endswith("OK"): + output = f"stdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" + pytest.fail(f"scenario {name!r} ({start_method}) failed (exit {proc.returncode})\n{output}") diff --git a/vortex-python/test/test_hf_datasets.py b/vortex-python/test/test_hf_datasets.py index 6990c907f85..5ae43a8bf8b 100644 --- a/vortex-python/test/test_hf_datasets.py +++ b/vortex-python/test/test_hf_datasets.py @@ -377,6 +377,30 @@ def test_materialize_filter_and_limit_combined(tmp_path: Path): assert dataset.to_list() == [{"text": "b", "label": 1}] +def test_materialize_filter_with_num_proc(tmp_path: Path): + """`num_proc` forks worker processes and pickles the filter into each one.""" + for shard in range(2): + write_vortex( + tmp_path / f"train-{shard}.vortex", + [{"text": f"{shard}-{i}", "label": i % 2} for i in range(50)], + ) + + dataset = vx_datasets.load_dataset( + tmp_path, + data_files="*.vortex", + split="train", + streaming=False, + filter=ve.column("label") == 1, + keep_in_memory=True, + num_proc=2, + ) + + assert isinstance(dataset, hf_datasets.Dataset) + rows = dataset.to_list() + assert len(rows) == 50 + assert {cast(int, row["label"]) for row in rows} == {1} + + def test_load_dataset_multi_split_without_mapping_raises(tmp_path: Path): write_vortex(tmp_path / "train.vortex", [{"text": "zero"}]) diff --git a/vortex-python/test/test_pickle.py b/vortex-python/test/test_pickle.py index 91f03c17f60..ce1fa84d929 100644 --- a/vortex-python/test/test_pickle.py +++ b/vortex-python/test/test_pickle.py @@ -3,9 +3,12 @@ import io import pickle +from pathlib import Path from typing import cast import pyarrow as pa +import pytest +from vortex.store import LocalStore import vortex as vx @@ -130,3 +133,24 @@ def test_pickle_file_api() -> None: restored = cast(vx.Array, pickle.Unpickler(file).load()) assert_arrow_array_roundtrip(arr, restored) + + +def test_pickle_vortex_file_reopens_by_path(tmp_path: Path) -> None: + path = tmp_path / "roundtrip.vortex" + vx.io.write(vx.array([1, 2, 3]), str(path)) + vxf = vx.open(str(path)) + + restored = cast(vx.VortexFile, pickle_roundtrip(vxf)) + + assert restored.path == str(path) + assert len(restored) == len(vxf) + assert restored.scan().read_all().to_arrow_array() == vxf.scan().read_all().to_arrow_array() + + +def test_pickle_vortex_file_rejects_explicit_store(tmp_path: Path) -> None: + path = tmp_path / "store.vortex" + vx.io.write(vx.array([1, 2, 3]), str(path)) + vxf = vx.open(str(path), store=LocalStore()) + + with pytest.raises(TypeError, match="object stores"): + _ = pickle.dumps(vxf)