From cbafed28842b9c52adeb8b1491a1eedbf1938f96 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 10 Aug 2026 02:04:14 +0100 Subject: [PATCH 1/7] Rebuild the Python runtime after fork `fork(2)` copies only the calling thread, so a forked child inherits a `CurrentThreadRuntime` whose worker threads no longer exist: the executor's sleeper list and the pool's handle list describe phantom threads, so `set_workers` believes it already has enough workers and spawns none, and every Vortex operation in the child blocks forever. The same applies to the process-global `blocking` pool behind `spawn_blocking`, whose inherited `idle_count` keeps it from growing. Tag the runtime with its owning pid and build a fresh one on first use from a different process, rather than trying to repair the inherited state; the stale state is leaked deliberately, since running its destructors would take locks that may not have survived the fork. An `os.register_at_fork(after_in_child=)` handler does the rebuild on the child's single-threaded startup path and warms the blocking pool. The shared session is repointed at the new executor, and the requested worker count is held outside the runtime so a child inherits it. `VortexFile` gains `path` and `__reduce__`, reopening by path in the receiving process, and the `datasets` filter path no longer needs its in-memory fallback now that a filter can be pickled into `num_proc` workers. Signed-off-by: Robert Kruszewski --- vortex-python/python/vortex/_lib/file.pyi | 6 + vortex-python/python/vortex/datasets.py | 56 +----- vortex-python/python/vortex/file.py | 5 + vortex-python/src/dataset.rs | 13 +- vortex-python/src/file.rs | 66 ++++++- vortex-python/src/io.rs | 9 +- vortex-python/src/lib.rs | 226 ++++++++++++++++++++-- vortex-python/src/runtime.rs | 30 +-- vortex-python/src/scan.rs | 8 +- vortex-python/src/session.rs | 29 ++- vortex-python/test/test_fork.py | 178 +++++++++++++++++ vortex-python/test/test_hf_datasets.py | 24 +++ vortex-python/test/test_pickle.py | 24 +++ 13 files changed, 573 insertions(+), 101 deletions(-) create mode 100644 vortex-python/test/test_fork.py diff --git a/vortex-python/python/vortex/_lib/file.pyi b/vortex-python/python/vortex/_lib/file.pyi index f8c928be459..a8732a43ec7 100644 --- a/vortex-python/python/vortex/_lib/file.pyi +++ b/vortex-python/python/vortex/_lib/file.pyi @@ -3,6 +3,8 @@ from typing import final +from typing_extensions import override + import polars as pl import pyarrow as pa @@ -22,6 +24,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..1d756529b42 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -28,12 +28,12 @@ use vortex::io::runtime::BlockingRuntime; 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?; @@ -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?; diff --git a/vortex-python/src/lib.rs b/vortex-python/src/lib.rs index d171afc1eb4..0b7b0dae3a7 100644 --- a/vortex-python/src/lib.rs +++ b/vortex-python/src/lib.rs @@ -2,12 +2,16 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::ops::Deref; -use std::sync::LazyLock; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; use log::LevelFilter; use pyo3::exceptions::PyRuntimeError; use pyo3::intern; use pyo3::prelude::*; +use pyo3::types::PyDict; use pyo3_log::Caching; use pyo3_log::Logger; @@ -36,28 +40,219 @@ mod serde; mod session; mod store; +use parking_lot::RwLock; +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); +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, tagged with the process that built +/// them. /// -/// 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") +/// `fork(2)` copies only the calling thread, so a forked child inherits a runtime whose worker +/// threads no longer exist: the executor's sleeper list and the pool's handle list both describe +/// phantom threads, and `CurrentThreadWorkerPool::set_workers` therefore believes it already has +/// enough workers and spawns none. Any Vortex operation in the child then blocks forever. Rather +/// than try to repair that state, we tag it with the owning pid and build a completely fresh +/// runtime the first time it is used from a different process. See [`with_state`]. +struct RuntimeState { + /// The process that built this runtime. + pid: u32, + runtime: CurrentThreadRuntime, + pool: CurrentThreadWorkerPool, +} + +impl RuntimeState { + fn new(pid: u32) -> Self { + let runtime = CurrentThreadRuntime::new(); + let pool = runtime.new_pool(); + pool.set_workers(requested_workers()); + Self { pid, runtime, pool } + } +} + +static RUNTIME_STATE: RwLock> = RwLock::new(None); + +/// The worker count to give a newly built pool, or [`WORKERS_UNSET`] to derive one. +/// +/// Held outside [`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; + } + 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); +} + +/// Runs `f` with this process's runtime state, building it if this is the first use or if the +/// process has forked since it was built. +/// +/// `f` runs under a read lock, so it must not itself call back into this module. +fn with_state(f: impl FnOnce(&RuntimeState) -> T) -> T { + let pid = std::process::id(); + loop { + { + let guard = RUNTIME_STATE.read(); + if let Some(state) = guard.as_ref() + && state.pid == pid + { + return f(state); + } + } + build_state(pid); + } +} + +fn build_state(pid: u32) { + let handle = { + let mut guard = RUNTIME_STATE.write(); + if guard.as_ref().is_some_and(|state| state.pid == pid) { + // Another thread got here first. + return; + } + let state = RuntimeState::new(pid); + let handle = state.runtime.handle(); + // Never run destructors on state inherited from a fork: dropping the pool takes its own + // mutex, and dropping the executor takes its task-slab mutex and then runs the destructors + // of every task it still holds. Any of those may have been held by a thread that did not + // survive the fork. Leaking is the lesser evil, and only ever happens once per fork. + if let Some(stale) = guard.replace(state) { + Box::leak(Box::new(stale)); + } + handle + }; + + // Repoint the shared session at the new executor. This must happen with the lock released, + // because building the session reads the runtime. + reset_session_handle(handle); +} + +/// The current-thread runtime backing Python Vortex operations. +pub(crate) fn current_runtime() -> CurrentThreadRuntime { + with_state(|state| 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 { + with_state(|state| f(&state.pool)) +} + +/// Rebuild the runtime in a freshly forked child process. +/// +/// Installed as an `os.register_at_fork(after_in_child=...)` handler. The pid check in +/// [`with_state`] alone would be enough, but running the rebuild from the fork handler keeps it +/// on the child's single-threaded startup path instead of racing whichever thread touches Vortex +/// first, and it also warms the global blocking-IO pool (see [`warm_blocking_pool`]). +fn reset_after_fork() { + // Only rebuild if this process's parent had actually built a runtime; otherwise leave it lazy + // so that children which never touch Vortex do not pay for worker threads. + if RUNTIME_STATE.read().is_none() { + return; } - pool -}); + warm_blocking_pool(¤t_runtime()); +} + +/// Force the process-global blocking-IO thread pool to spawn at least one live thread. +/// +/// Vortex reads route through `Handle::spawn_blocking`, which is backed by the `blocking` crate's +/// process-global pool. That pool only grows while `queue.len() > idle_count * 5`, and a forked +/// child inherits an `idle_count` counting threads that no longer exist — so the child's first reads +/// are queued and never run. Submitting no-op tasks pushes the queue past the inherited threshold, +/// which spawns a real thread; that thread then drains the queue and serves subsequent work. +/// +/// Probes go out in small batches with a short wait in between: batching keeps the number of probes +/// needed (`5 * inherited_idle_count`) from costing one wait each, while a small batch size keeps the +/// pool from overshooting into dozens of threads before the first one starts draining the queue. +/// +/// Completion is signalled through an atomic polled with `thread::sleep`, not a channel. This runs +/// inside an `atfork` child handler, where `std::sync::mpsc` must not be used: its blocking receive +/// parks on `std::thread::park`, which on macOS is backed by a libdispatch semaphore, and libdispatch +/// traps (`SIGTRAP`) when a semaphore created before the fork is waited on in the child. +fn warm_blocking_pool(runtime: &CurrentThreadRuntime) { + /// Bounds the warmup in case the pool is unrecoverable, e.g. its mutex was held across the fork. + const MAX_PROBES: usize = 8192; + /// Probes per completion check. + const BATCH: usize = 8; + /// Wait between completion checks. A healthy pool answers the first batch within this window, so + /// the common case costs one short sleep. + const POLL: Duration = Duration::from_millis(2); + /// Extra checks after the last batch, for a pool that is slow rather than broken. + const GRACE_POLLS: usize = 250; + + let handle = runtime.handle(); + let completed = Arc::new(AtomicUsize::new(0)); + + let mut submitted = 0; + while submitted < MAX_PROBES { + for _ in 0..BATCH { + let completed = Arc::clone(&completed); + handle + .spawn_blocking(move || { + completed.fetch_add(1, Ordering::Release); + }) + .detach(); + } + submitted += BATCH; + std::thread::sleep(POLL); + if completed.load(Ordering::Acquire) > 0 { + return; + } + } + + for _ in 0..GRACE_POLLS { + std::thread::sleep(POLL); + if completed.load(Ordering::Acquire) > 0 { + return; + } + } + + log::warn!( + "Vortex could not restart its blocking IO thread pool after fork; reads may not complete" + ); +} + +/// Install the `after_in_child` fork handler that rebuilds the runtime. +fn register_at_fork(py: Python) -> PyResult<()> { + let handler = wrap_pyfunction!(_reset_after_fork, py)?; + let kwargs = PyDict::new(py); + kwargs.set_item(intern!(py, "after_in_child"), handler)?; + py.import("os")? + .call_method(intern!(py, "register_at_fork"), (), Some(&kwargs))?; + Ok(()) +} + +/// Python entry point for [`reset_after_fork`]. +#[pyfunction] +fn _reset_after_fork(py: Python) { + py.detach(reset_after_fork); +} /// Vortex is an Apache Arrow-compatible toolkit for working with compressed array data. #[cfg(feature = "extension-module")] @@ -71,6 +266,9 @@ fn _lib(py: Python, m: &Bound) -> PyResult<()> { .map_err(|err| PyRuntimeError::new_err(format!("could not initialize logger {err}"))) })?; + // `fork(2)` leaves the inherited Vortex runtime unusable, so rebuild it in the child. + register_at_fork(py)?; + // Initialize our submodules, living under vortex._lib arrays::init(py, m)?; #[cfg(feature = "tui")] 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..19bf9818645 100644 --- a/vortex-python/src/session.rs +++ b/vortex-python/src/session.rs @@ -1,18 +1,37 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::LazyLock; +use std::sync::OnceLock; use vortex::VortexSessionDefault; use vortex::io::runtime::BlockingRuntime; +use vortex::io::runtime::Handle; use vortex::io::session::RuntimeSessionExt; +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())); +/// A `OnceLock` rather than a `LazyLock` because initialization is re-entrant: building the session +/// builds the runtime, and building the runtime calls [`reset_session_handle`], which must be able +/// to observe "not initialized yet" without blocking. +static SESSION: OnceLock = OnceLock::new(); pub(crate) fn session() -> &'static VortexSession { - &SESSION + SESSION.get_or_init(|| 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. +pub(crate) fn reset_session_handle(handle: Handle) { + if let Some(session) = SESSION.get() { + session.session().with_handle(handle); + } } diff --git a/vortex-python/test/test_fork.py b/vortex-python/test/test_fork.py new file mode 100644 index 00000000000..4f7c2ca334d --- /dev/null +++ b/vortex-python/test/test_fork.py @@ -0,0 +1,178 @@ +# 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. + +Each scenario runs in a **fresh interpreter** that imports nothing but Vortex, and only then forks. +That keeps the tests measuring Vortex's own fork 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 subprocess +import sys +import textwrap +from pathlib import Path + +import pyarrow as pa +import pytest + +import vortex as vx + +pytestmark = pytest.mark.skipif(sys.platform.startswith("win"), reason="fork(2) is not available on Windows") + +TIMEOUT_SECONDS = 120 +ROWS = 10_000 + +PREAMBLE = textwrap.dedent( + """ + import multiprocessing, sys + import vortex as vx + import vortex.expr as ve + + PATH = sys.argv[1] + CTX = multiprocessing.get_context("fork") + + def in_fork(target, *args): + "Run `target` in a forked child and return its result, or raise on hang/error." + parent_conn, child_conn = CTX.Pipe(duplex=False) + + def entry(): + try: + child_conn.send(("ok", target(*args))) + except BaseException as err: + child_conn.send(("err", f"{type(err).__name__}: {err}")) + finally: + child_conn.close() + + process = CTX.Process(target=entry) + process.start() + try: + if not parent_conn.poll(60): + process.terminate() + raise AssertionError("forked child produced no result within 60s") + status, payload = parent_conn.recv() + finally: + process.join(30) + if status == "err": + raise AssertionError(f"forked child raised {payload}") + return payload + """ +) + +# Each scenario must print nothing but a trailing "OK" on success. +SCENARIOS: dict[str, str] = { + # The parent has already built a runtime, so the child inherits a populated (and useless) one + # and must build its own before it can read anything. + "read_after_parent_read": """ + assert len(vx.open(PATH).scan().read_all()) == ROWS + + def child(): + return len(vx.open(PATH).scan().read_all()) + + assert in_fork(child) == ROWS + print("OK") + """, + # The child must end up with real worker threads, not the parent's phantom handles. + "child_has_live_workers": """ + assert len(vx.open(PATH).scan().read_all()) == ROWS + assert in_fork(vx.worker_threads) >= 1 + print("OK") + """, + # An Expr filter evaluated in the child. + "expr_filter_in_child": """ + assert len(vx.open(PATH).scan().read_all()) == ROWS + expr = (ve.column("age") >= ROWS - 2) & ve.like(ve.column("name"), "person-%") + + def child(expr): + table = vx.open(PATH).scan(["name"], expr=expr).read_all().to_arrow_table() + return table.column("name").to_pylist() + + assert in_fork(child, expr) == [f"person-{ROWS - 2}", f"person-{ROWS - 1}"] + print("OK") + """, + # The worker count configured in the parent must carry over to the child's fresh pool. + "child_inherits_worker_count": """ + vx.set_worker_threads(3) + assert len(vx.open(PATH).scan().read_all()) == ROWS + assert in_fork(vx.worker_threads) == 3 + print("OK") + """, + # A fork-backed Pool pickles both the callable and its arguments: the shape `datasets` uses for + # `num_proc`, and the reason Expr needs `__reduce__`. + "pool_map_with_pickled_expr": """ + assert len(vx.open(PATH).scan().read_all()) == 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=60) + except multiprocessing.TimeoutError: + pool.terminate() + raise AssertionError("fork-backed Pool did not finish within 60s") + assert counts == [100, 100], counts + print("OK") + """, + # A pickled VortexFile is reopened by path in the child. + "pickled_file_reads_in_child": """ + import pickle + vxf = vx.open(PATH) + assert len(vxf.scan().read_all()) == ROWS + restored = pickle.loads(pickle.dumps(vxf)) + assert restored.path == PATH + assert in_fork(read_file, restored) == ROWS + print("OK") + """, +} + +# Top-level so `Pool.map` can pickle it by reference. +HELPERS = textwrap.dedent( + """ + def count_filtered(args): + path, expr = args + return len(vx.open(path).scan(expr=expr).read_all()) + + def read_file(file): + return len(file.scan().read_all()) + """ +) + + +@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", sorted(SCENARIOS)) +def test_fork_scenario(name: str, people: Path) -> None: + script = "\n".join( + [ + PREAMBLE, + f"ROWS = {ROWS}", + HELPERS, + "if __name__ == '__main__':", + textwrap.indent(textwrap.dedent(SCENARIOS[name]), " "), + ] + ) + + try: + proc = subprocess.run( + [sys.executable, "-c", script, str(people)], + capture_output=True, + text=True, + timeout=TIMEOUT_SECONDS, + check=False, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"scenario {name!r} did not finish within {TIMEOUT_SECONDS}s (likely a fork deadlock)") + + if proc.returncode != 0 or not proc.stdout.strip().endswith("OK"): + pytest.fail( + f"scenario {name!r} failed (exit {proc.returncode})\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" + ) 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) From b3f076797862693ee21d77f0f4900cceeaef9fae Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 12 Aug 2026 13:33:32 +0100 Subject: [PATCH 2/7] less Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 - vortex-io/src/runtime/blocking_pool.rs | 236 +++++++++++++++++++++++++ vortex-io/src/runtime/current.rs | 9 +- vortex-io/src/runtime/mod.rs | 2 + vortex-io/src/runtime/pool.rs | 6 +- vortex-io/src/runtime/smol.rs | 42 ++++- vortex-io/src/std_file/mod.rs | 2 + vortex-io/src/std_file/write.rs | 83 +++++++++ vortex-python/Cargo.toml | 1 - vortex-python/src/io.rs | 6 +- vortex-python/src/lib.rs | 216 +++++++--------------- vortex-python/src/session.rs | 64 ++++++- vortex-python/test/test_fork.py | 28 +++ 13 files changed, 524 insertions(+), 172 deletions(-) create mode 100644 vortex-io/src/runtime/blocking_pool.rs create mode 100644 vortex-io/src/std_file/write.rs diff --git a/Cargo.lock b/Cargo.lock index 9a7274a774b..0b5e7881627 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10464,7 +10464,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/vortex-io/src/runtime/blocking_pool.rs b/vortex-io/src/runtime/blocking_pool.rs new file mode 100644 index 00000000000..b44ca948f7f --- /dev/null +++ b/vortex-io/src/runtime/blocking_pool.rs @@ -0,0 +1,236 @@ +// 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 kanal::Receiver; +use kanal::Sender; +use kanal::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)); + if self + .sender + .send(Job { + cancelled: Arc::clone(&cancelled), + task, + }) + .is_err() + { + vortex_panic!("cannot spawn blocking work on a shut down runtime"); + } + self.grow(); + Box::new(BlockingAbortHandle { cancelled }) + } + + fn grow(&self) { + let mut state = self.state.lock(); + // An unbounded kanal send hands work directly to a waiting receiver. A queued job therefore + // means every existing worker is busy (or transitioning out of an idle wait), so add one + // worker for this submission while capacity remains. + if self.sender.is_empty() || 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}"); + } + } +} + +impl Drop for BlockingPool { + fn drop(&mut self) { + drop(self.sender.close()); + } +} + +#[derive(Default)] +struct PoolState { + thread_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 { + match receiver.recv_timeout(KEEP_ALIVE) { + Ok(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()))); + } + Err(kanal::ReceiveErrorTimeout::Timeout) => { + let mut state = state.lock(); + if receiver.is_empty() { + state.thread_count -= 1; + return; + } + } + Err(kanal::ReceiveErrorTimeout::Closed | kanal::ReceiveErrorTimeout::SendClosed) => { + let mut state = state.lock(); + state.thread_count -= 1; + return; + } + } + } +} + +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 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(()) + } +} 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/src/io.rs b/vortex-python/src/io.rs index 1d756529b42..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,6 +24,7 @@ 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; @@ -261,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()) @@ -397,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 0b7b0dae3a7..0183bd772d8 100644 --- a/vortex-python/src/lib.rs +++ b/vortex-python/src/lib.rs @@ -2,16 +2,21 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::ops::Deref; -use std::sync::Arc; +#[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 std::time::Duration; use log::LevelFilter; use pyo3::exceptions::PyRuntimeError; use pyo3::intern; use pyo3::prelude::*; -use pyo3::types::PyDict; use pyo3_log::Caching; use pyo3_log::Logger; @@ -40,45 +45,49 @@ mod serde; mod session; mod store; -use parking_lot::RwLock; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; use vortex::io::runtime::current::CurrentThreadWorkerPool; use vortex::utils::parallelism::get_available_parallelism; +#[cfg(unix)] use crate::session::reset_session_handle; -/// The current-thread runtime and its background worker pool, tagged with the process that built -/// them. -/// -/// `fork(2)` copies only the calling thread, so a forked child inherits a runtime whose worker -/// threads no longer exist: the executor's sleeper list and the pool's handle list both describe -/// phantom threads, and `CurrentThreadWorkerPool::set_workers` therefore believes it already has -/// enough workers and spawns none. Any Vortex operation in the child then blocks forever. Rather -/// than try to repair that state, we tag it with the owning pid and build a completely fresh -/// runtime the first time it is used from a different process. See [`with_state`]. +/// The current-thread runtime and its background worker pool. struct RuntimeState { - /// The process that built this runtime. - pid: u32, runtime: CurrentThreadRuntime, pool: CurrentThreadWorkerPool, } impl RuntimeState { - fn new(pid: u32) -> Self { + fn new() -> Self { let runtime = CurrentThreadRuntime::new(); let pool = runtime.new_pool(); pool.set_workers(requested_workers()); - Self { pid, runtime, pool } + Self { runtime, pool } } } -static RUNTIME_STATE: RwLock> = RwLock::new(None); +#[cfg(not(unix))] +static RUNTIME_STATE: LazyLock = LazyLock::new(RuntimeState::new); + +/// A process-tagged initialization cell. +/// +/// 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 [`RUNTIME_STATE`] so that a forked child inherits the parent's configuration even -/// though it discards the parent's runtime. +/// 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; @@ -107,52 +116,56 @@ pub(crate) fn set_requested_workers(workers: usize) { REQUESTED_WORKERS.store(workers, Ordering::Relaxed); } -/// Runs `f` with this process's runtime state, building it if this is the first use or if the -/// process has forked since it was built. -/// -/// `f` runs under a read lock, so it must not itself call back into this module. -fn with_state(f: impl FnOnce(&RuntimeState) -> T) -> T { +#[cfg(unix)] +fn runtime_slot() -> &'static RuntimeSlot { let pid = std::process::id(); loop { - { - let guard = RUNTIME_STATE.read(); - if let Some(state) = guard.as_ref() - && state.pid == pid - { - return f(state); + 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) }); } } - build_state(pid); } } -fn build_state(pid: u32) { - let handle = { - let mut guard = RUNTIME_STATE.write(); - if guard.as_ref().is_some_and(|state| state.pid == pid) { - // Another thread got here first. - return; - } - let state = RuntimeState::new(pid); - let handle = state.runtime.handle(); - // Never run destructors on state inherited from a fork: dropping the pool takes its own - // mutex, and dropping the executor takes its task-slab mutex and then runs the destructors - // of every task it still holds. Any of those may have been held by a thread that did not - // survive the fork. Leaking is the lesser evil, and only ever happens once per fork. - if let Some(stale) = guard.replace(state) { - Box::leak(Box::new(stale)); - } - handle - }; +#[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 + }) +} - // Repoint the shared session at the new executor. This must happen with the lock released, - // because building the session reads the runtime. - reset_session_handle(handle); +#[cfg(not(unix))] +fn runtime_state() -> &'static RuntimeState { + &RUNTIME_STATE } /// The current-thread runtime backing Python Vortex operations. pub(crate) fn current_runtime() -> CurrentThreadRuntime { - with_state(|state| state.runtime.clone()) + runtime_state().runtime.clone() } /// Runs `f` with the worker pool that drives [`current_runtime`]'s executor in the background. @@ -161,97 +174,7 @@ pub(crate) fn current_runtime() -> CurrentThreadRuntime { /// 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 { - with_state(|state| f(&state.pool)) -} - -/// Rebuild the runtime in a freshly forked child process. -/// -/// Installed as an `os.register_at_fork(after_in_child=...)` handler. The pid check in -/// [`with_state`] alone would be enough, but running the rebuild from the fork handler keeps it -/// on the child's single-threaded startup path instead of racing whichever thread touches Vortex -/// first, and it also warms the global blocking-IO pool (see [`warm_blocking_pool`]). -fn reset_after_fork() { - // Only rebuild if this process's parent had actually built a runtime; otherwise leave it lazy - // so that children which never touch Vortex do not pay for worker threads. - if RUNTIME_STATE.read().is_none() { - return; - } - warm_blocking_pool(¤t_runtime()); -} - -/// Force the process-global blocking-IO thread pool to spawn at least one live thread. -/// -/// Vortex reads route through `Handle::spawn_blocking`, which is backed by the `blocking` crate's -/// process-global pool. That pool only grows while `queue.len() > idle_count * 5`, and a forked -/// child inherits an `idle_count` counting threads that no longer exist — so the child's first reads -/// are queued and never run. Submitting no-op tasks pushes the queue past the inherited threshold, -/// which spawns a real thread; that thread then drains the queue and serves subsequent work. -/// -/// Probes go out in small batches with a short wait in between: batching keeps the number of probes -/// needed (`5 * inherited_idle_count`) from costing one wait each, while a small batch size keeps the -/// pool from overshooting into dozens of threads before the first one starts draining the queue. -/// -/// Completion is signalled through an atomic polled with `thread::sleep`, not a channel. This runs -/// inside an `atfork` child handler, where `std::sync::mpsc` must not be used: its blocking receive -/// parks on `std::thread::park`, which on macOS is backed by a libdispatch semaphore, and libdispatch -/// traps (`SIGTRAP`) when a semaphore created before the fork is waited on in the child. -fn warm_blocking_pool(runtime: &CurrentThreadRuntime) { - /// Bounds the warmup in case the pool is unrecoverable, e.g. its mutex was held across the fork. - const MAX_PROBES: usize = 8192; - /// Probes per completion check. - const BATCH: usize = 8; - /// Wait between completion checks. A healthy pool answers the first batch within this window, so - /// the common case costs one short sleep. - const POLL: Duration = Duration::from_millis(2); - /// Extra checks after the last batch, for a pool that is slow rather than broken. - const GRACE_POLLS: usize = 250; - - let handle = runtime.handle(); - let completed = Arc::new(AtomicUsize::new(0)); - - let mut submitted = 0; - while submitted < MAX_PROBES { - for _ in 0..BATCH { - let completed = Arc::clone(&completed); - handle - .spawn_blocking(move || { - completed.fetch_add(1, Ordering::Release); - }) - .detach(); - } - submitted += BATCH; - std::thread::sleep(POLL); - if completed.load(Ordering::Acquire) > 0 { - return; - } - } - - for _ in 0..GRACE_POLLS { - std::thread::sleep(POLL); - if completed.load(Ordering::Acquire) > 0 { - return; - } - } - - log::warn!( - "Vortex could not restart its blocking IO thread pool after fork; reads may not complete" - ); -} - -/// Install the `after_in_child` fork handler that rebuilds the runtime. -fn register_at_fork(py: Python) -> PyResult<()> { - let handler = wrap_pyfunction!(_reset_after_fork, py)?; - let kwargs = PyDict::new(py); - kwargs.set_item(intern!(py, "after_in_child"), handler)?; - py.import("os")? - .call_method(intern!(py, "register_at_fork"), (), Some(&kwargs))?; - Ok(()) -} - -/// Python entry point for [`reset_after_fork`]. -#[pyfunction] -fn _reset_after_fork(py: Python) { - py.detach(reset_after_fork); + f(&runtime_state().pool) } /// Vortex is an Apache Arrow-compatible toolkit for working with compressed array data. @@ -266,9 +189,6 @@ fn _lib(py: Python, m: &Bound) -> PyResult<()> { .map_err(|err| PyRuntimeError::new_err(format!("could not initialize logger {err}"))) })?; - // `fork(2)` leaves the inherited Vortex runtime unusable, so rebuild it in the child. - register_at_fork(py)?; - // Initialize our submodules, living under vortex._lib arrays::init(py, m)?; #[cfg(feature = "tui")] diff --git a/vortex-python/src/session.rs b/vortex-python/src/session.rs index 19bf9818645..651addb5457 100644 --- a/vortex-python/src/session.rs +++ b/vortex-python/src/session.rs @@ -1,24 +1,70 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::OnceLock; +#[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::current_runtime; -/// A `OnceLock` rather than a `LazyLock` because initialization is re-entrant: building the session -/// builds the runtime, and building the runtime calls [`reset_session_handle`], which must be able -/// to observe "not initialized yet" without blocking. -static SESSION: OnceLock = OnceLock::new(); +#[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 { - SESSION.get_or_init(|| VortexSession::default().with_handle(current_runtime().handle())) + // 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. @@ -30,8 +76,12 @@ pub(crate) fn session() -> &'static VortexSession { /// /// 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) { - if let Some(session) = SESSION.get() { + 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/test_fork.py b/vortex-python/test/test_fork.py index 4f7c2ca334d..03fc8c71858 100644 --- a/vortex-python/test/test_fork.py +++ b/vortex-python/test/test_fork.py @@ -77,6 +77,34 @@ def child(): assert in_fork(child) == ROWS print("OK") """, + # 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. + "read_after_saturated_blocking_pool": """ + import os + os.environ["BLOCKING_MAX_THREADS"] = "1" + assert len(vx.open(PATH).scan().read_all()) == ROWS + + def child(): + return len(vx.open(PATH).scan().read_all()) + + assert in_fork(child) == ROWS + print("OK") + """, + # 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. + "write_after_parent_write": """ + import pyarrow as pa + values = vx.array(pa.array([{"value": i} for i in range(10)])) + vx.io.write(values, f"{PATH}.parent-write") + + def child(): + child_path = f"{PATH}.child-write" + vx.io.write(values, child_path) + return len(vx.open(child_path)) + + assert in_fork(child) == 10 + print("OK") + """, # The child must end up with real worker threads, not the parent's phantom handles. "child_has_live_workers": """ assert len(vx.open(PATH).scan().read_all()) == ROWS From 6532dbbca5c607843acbae19b7ed8311c380b5eb Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 12 Aug 2026 13:39:45 +0100 Subject: [PATCH 3/7] format Signed-off-by: Robert Kruszewski --- vortex-python/python/vortex/_lib/file.pyi | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vortex-python/python/vortex/_lib/file.pyi b/vortex-python/python/vortex/_lib/file.pyi index a8732a43ec7..901e0323eec 100644 --- a/vortex-python/python/vortex/_lib/file.pyi +++ b/vortex-python/python/vortex/_lib/file.pyi @@ -3,10 +3,9 @@ from typing import final -from typing_extensions import override - import polars as pl import pyarrow as pa +from typing_extensions import override from vortex.type_aliases import IntoProjection From 9ae4715b0ab93c3dc3eb09fa1c778e38c0ff239c Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 12 Aug 2026 14:56:33 +0100 Subject: [PATCH 4/7] supressions Signed-off-by: Robert Kruszewski --- vortex-ffi/tsan_suppressions.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vortex-ffi/tsan_suppressions.txt b/vortex-ffi/tsan_suppressions.txt index 6a2b65e351e..28128686968 100644 --- a/vortex-ffi/tsan_suppressions.txt +++ b/vortex-ffi/tsan_suppressions.txt @@ -6,6 +6,8 @@ # where ordering is correct but uses an explicit fence and not relaxed load race:oneshot-*/src/channel.rs race:oneshot-*/src/lib.rs +race:kanal-*/src/lib.rs +race:kanal-*/src/signal.rs # https://github.com/google/sanitizers/issues/1415 # crossbeam-epoch's epoch-based reclamation synchronizes through SeqCst fences, From 53d89adbf2b46c605855499cd5ee97ab1569e765 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 12 Aug 2026 16:10:42 +0100 Subject: [PATCH 5/7] more Signed-off-by: Robert Kruszewski --- vortex-ffi/tsan_suppressions.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vortex-ffi/tsan_suppressions.txt b/vortex-ffi/tsan_suppressions.txt index 28128686968..8a9c7113dcf 100644 --- a/vortex-ffi/tsan_suppressions.txt +++ b/vortex-ffi/tsan_suppressions.txt @@ -8,6 +8,10 @@ race:oneshot-*/src/channel.rs race:oneshot-*/src/lib.rs race:kanal-*/src/lib.rs race:kanal-*/src/signal.rs +# Suppressing the Kanal frames does not make TSan propagate the handoff's happens-before edge. +# Catch the resulting reports at the pool's worker boundary instead of suppressing each operation +# that a blocking job may perform. +race:vortex_io*blocking_pool*worker_loop # https://github.com/google/sanitizers/issues/1415 # crossbeam-epoch's epoch-based reclamation synchronizes through SeqCst fences, From cf98f9c06b3e7e2b8cc3dfec2339252dac4f0aee Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 13 Aug 2026 13:48:01 +0100 Subject: [PATCH 6/7] replace kanal with crossbeam Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 + Cargo.toml | 1 + vortex-ffi/tsan_suppressions.txt | 6 -- vortex-io/Cargo.toml | 1 + vortex-io/src/runtime/blocking_pool.rs | 132 +++++++++++++++++++------ 5 files changed, 107 insertions(+), 34 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0b5e7881627..fee1e3ebd34 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10235,6 +10235,7 @@ dependencies = [ "async-stream", "async-trait", "bytes", + "crossbeam-channel", "custom-labels", "futures", "glob", 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-ffi/tsan_suppressions.txt b/vortex-ffi/tsan_suppressions.txt index 8a9c7113dcf..6a2b65e351e 100644 --- a/vortex-ffi/tsan_suppressions.txt +++ b/vortex-ffi/tsan_suppressions.txt @@ -6,12 +6,6 @@ # where ordering is correct but uses an explicit fence and not relaxed load race:oneshot-*/src/channel.rs race:oneshot-*/src/lib.rs -race:kanal-*/src/lib.rs -race:kanal-*/src/signal.rs -# Suppressing the Kanal frames does not make TSan propagate the handoff's happens-before edge. -# Catch the resulting reports at the pool's worker boundary instead of suppressing each operation -# that a blocking job may perform. -race:vortex_io*blocking_pool*worker_loop # https://github.com/google/sanitizers/issues/1415 # crossbeam-epoch's epoch-based reclamation synchronizes through SeqCst fences, 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 index b44ca948f7f..fc315844831 100644 --- a/vortex-io/src/runtime/blocking_pool.rs +++ b/vortex-io/src/runtime/blocking_pool.rs @@ -8,9 +8,11 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; -use kanal::Receiver; -use kanal::Sender; -use kanal::unbounded; +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; @@ -50,6 +52,7 @@ impl BlockingPool { 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 { @@ -60,16 +63,15 @@ impl BlockingPool { { vortex_panic!("cannot spawn blocking work on a shut down runtime"); } - self.grow(); + 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) { - let mut state = self.state.lock(); - // An unbounded kanal send hands work directly to a waiting receiver. A queued job therefore - // means every existing worker is busy (or transitioning out of an idle wait), so add one - // worker for this submission while capacity remains. - if self.sender.is_empty() || state.thread_count >= self.thread_limit { + fn grow(&self, state: &mut PoolState) { + if state.thread_count >= self.thread_limit { return; } state.thread_count += 1; @@ -86,15 +88,13 @@ impl BlockingPool { } } -impl Drop for BlockingPool { - fn drop(&mut self) { - drop(self.sender.close()); - } -} - #[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 { @@ -122,29 +122,65 @@ impl AbortHandle for BlockingAbortHandle { 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) => { - // `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()))); + 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(kanal::ReceiveErrorTimeout::Timeout) => { - let mut state = state.lock(); - if receiver.is_empty() { - state.thread_count -= 1; - return; + 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(kanal::ReceiveErrorTimeout::Closed | kanal::ReceiveErrorTimeout::SendClosed) => { - let mut state = state.lock(); - state.thread_count -= 1; + 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() @@ -157,6 +193,7 @@ fn max_threads() -> usize { mod tests { use std::sync::Barrier; use std::sync::mpsc; + use std::time::Instant; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -233,4 +270,43 @@ mod tests { 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(()) + } } From 5f7a575bf77776ac60ef2697425f5b32ff6007c7 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 13 Aug 2026 16:39:05 +0100 Subject: [PATCH 7/7] tests Signed-off-by: Robert Kruszewski --- vortex-python/test/child_scenarios.py | 190 +++++++++++++++++++++++++ vortex-python/test/test_fork.py | 191 ++++---------------------- 2 files changed, 217 insertions(+), 164 deletions(-) create mode 100644 vortex-python/test/child_scenarios.py 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 index 03fc8c71858..832ad6493af 100644 --- a/vortex-python/test/test_fork.py +++ b/vortex-python/test/test_fork.py @@ -3,19 +3,25 @@ """The Vortex runtime must be rebuilt in forked children, or every read there hangs. -Each scenario runs in a **fresh interpreter** that imports nothing but Vortex, and only then forks. -That keeps the tests measuring Vortex's own fork 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. +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 -import textwrap from pathlib import Path import pyarrow as pa @@ -23,151 +29,16 @@ import vortex as vx -pytestmark = pytest.mark.skipif(sys.platform.startswith("win"), reason="fork(2) is not available on Windows") +from .child_scenarios import ROWS, SCENARIOS TIMEOUT_SECONDS = 120 -ROWS = 10_000 - -PREAMBLE = textwrap.dedent( - """ - import multiprocessing, sys - import vortex as vx - import vortex.expr as ve - - PATH = sys.argv[1] - CTX = multiprocessing.get_context("fork") - - def in_fork(target, *args): - "Run `target` in a forked child and return its result, or raise on hang/error." - parent_conn, child_conn = CTX.Pipe(duplex=False) - - def entry(): - try: - child_conn.send(("ok", target(*args))) - except BaseException as err: - child_conn.send(("err", f"{type(err).__name__}: {err}")) - finally: - child_conn.close() - - process = CTX.Process(target=entry) - process.start() - try: - if not parent_conn.poll(60): - process.terminate() - raise AssertionError("forked child produced no result within 60s") - status, payload = parent_conn.recv() - finally: - process.join(30) - if status == "err": - raise AssertionError(f"forked child raised {payload}") - return payload - """ -) - -# Each scenario must print nothing but a trailing "OK" on success. -SCENARIOS: dict[str, str] = { - # The parent has already built a runtime, so the child inherits a populated (and useless) one - # and must build its own before it can read anything. - "read_after_parent_read": """ - assert len(vx.open(PATH).scan().read_all()) == ROWS - - def child(): - return len(vx.open(PATH).scan().read_all()) - - assert in_fork(child) == ROWS - print("OK") - """, - # 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. - "read_after_saturated_blocking_pool": """ - import os - os.environ["BLOCKING_MAX_THREADS"] = "1" - assert len(vx.open(PATH).scan().read_all()) == ROWS - - def child(): - return len(vx.open(PATH).scan().read_all()) - - assert in_fork(child) == ROWS - print("OK") - """, - # 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. - "write_after_parent_write": """ - import pyarrow as pa - values = vx.array(pa.array([{"value": i} for i in range(10)])) - vx.io.write(values, f"{PATH}.parent-write") - - def child(): - child_path = f"{PATH}.child-write" - vx.io.write(values, child_path) - return len(vx.open(child_path)) - - assert in_fork(child) == 10 - print("OK") - """, - # The child must end up with real worker threads, not the parent's phantom handles. - "child_has_live_workers": """ - assert len(vx.open(PATH).scan().read_all()) == ROWS - assert in_fork(vx.worker_threads) >= 1 - print("OK") - """, - # An Expr filter evaluated in the child. - "expr_filter_in_child": """ - assert len(vx.open(PATH).scan().read_all()) == ROWS - expr = (ve.column("age") >= ROWS - 2) & ve.like(ve.column("name"), "person-%") - - def child(expr): - table = vx.open(PATH).scan(["name"], expr=expr).read_all().to_arrow_table() - return table.column("name").to_pylist() - - assert in_fork(child, expr) == [f"person-{ROWS - 2}", f"person-{ROWS - 1}"] - print("OK") - """, - # The worker count configured in the parent must carry over to the child's fresh pool. - "child_inherits_worker_count": """ - vx.set_worker_threads(3) - assert len(vx.open(PATH).scan().read_all()) == ROWS - assert in_fork(vx.worker_threads) == 3 - print("OK") - """, - # A fork-backed Pool pickles both the callable and its arguments: the shape `datasets` uses for - # `num_proc`, and the reason Expr needs `__reduce__`. - "pool_map_with_pickled_expr": """ - assert len(vx.open(PATH).scan().read_all()) == 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=60) - except multiprocessing.TimeoutError: - pool.terminate() - raise AssertionError("fork-backed Pool did not finish within 60s") - assert counts == [100, 100], counts - print("OK") - """, - # A pickled VortexFile is reopened by path in the child. - "pickled_file_reads_in_child": """ - import pickle - vxf = vx.open(PATH) - assert len(vxf.scan().read_all()) == ROWS - restored = pickle.loads(pickle.dumps(vxf)) - assert restored.path == PATH - assert in_fork(read_file, restored) == ROWS - print("OK") - """, -} - -# Top-level so `Pool.map` can pickle it by reference. -HELPERS = textwrap.dedent( - """ - def count_filtered(args): - path, expr = args - return len(vx.open(path).scan(expr=expr).read_all()) - - def read_file(file): - return len(file.scan().read_all()) - """ -) +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") @@ -177,30 +48,22 @@ def people(tmp_path_factory: pytest.TempPathFactory) -> Path: return path -@pytest.mark.parametrize("name", sorted(SCENARIOS)) -def test_fork_scenario(name: str, people: Path) -> None: - script = "\n".join( - [ - PREAMBLE, - f"ROWS = {ROWS}", - HELPERS, - "if __name__ == '__main__':", - textwrap.indent(textwrap.dedent(SCENARIOS[name]), " "), - ] - ) +@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, "-c", script, str(people)], + [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} did not finish within {TIMEOUT_SECONDS}s (likely a fork deadlock)") + 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"): - pytest.fail( - f"scenario {name!r} failed (exit {proc.returncode})\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" - ) + output = f"stdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" + pytest.fail(f"scenario {name!r} ({start_method}) failed (exit {proc.returncode})\n{output}")