From 5fd84e9b24d01dba1cd5f9ee32a1a04c2a681256 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Mon, 22 Jun 2026 14:34:53 +0000 Subject: [PATCH 1/3] feat(python): package CUDA as an optional extension Keep the base vortex-data wheel CPU-only and model CUDA support as a separate vortex-data-cuda extension package imported as vortex_cuda. The base package exposes vortex.cuda_extension_installed() to check whether the optional extension is importable, while vortex_cuda.cuda_available() performs the runtime CUDA driver/device probe. Wire vortex-data[cuda] to depend on the exact matching vortex-data-cuda version, and have the CUDA package depend on the exact matching vortex-data version. Release automation updates those pins with the workspace version, but this change intentionally does not add CUDA wheel publishing jobs yet. Keep vortex-python-cuda out of the uv workspace so uv --all-packages on non-GPU CI does not build the CUDA crate chain. Exercise it through the GPU CUDA workflow instead, and exclude the thin PyO3 shim from broad workspace doc/doctest/instrumented jobs where it has no Rust tests or doctests. Embed generated PTX into vortex-cuda when nvcc is available so installed wheels do not depend on build-machine kernel file paths at runtime. Runtime kernel loading is embedded-only; the old VORTEX_CUDA_KERNELS_DIR path override and Codspeed kernel staging were removed. Make cuda_available() resilient on driverless hosts by probing libcuda with cudarc before creating a CUDA context, so CPU-only machines return false instead of panickingq Signed-off-by: Alexander Droste --- .github/workflows/ci.yml | 8 +-- .github/workflows/codspeed.yml | 5 -- .github/workflows/cuda.yaml | 22 ++++++- .github/workflows/musl.yml | 2 +- .github/workflows/package.yml | 24 +++++++ .github/workflows/rust-instrumented.yml | 2 +- Cargo.lock | 8 +++ Cargo.toml | 1 + docs/api/python/runtime.rst | 3 + pyproject.toml | 5 ++ uv.lock | 16 ++++- vortex-cuda/build.rs | 63 ++++++++++++++++--- vortex-cuda/src/kernel/mod.rs | 51 +++++---------- vortex-cuda/src/lib.rs | 30 ++++++--- vortex-python-cuda/Cargo.toml | 31 +++++++++ vortex-python-cuda/pyproject.toml | 34 ++++++++++ .../python/vortex_cuda/__init__.py | 6 ++ .../python/vortex_cuda/_lib.pyi | 4 ++ .../python/vortex_cuda/py.typed | 0 vortex-python-cuda/src/lib.rs | 28 +++++++++ vortex-python-cuda/test/test_cuda.py | 29 +++++++++ vortex-python/pyproject.toml | 1 + vortex-python/python/vortex/__init__.py | 19 ++++++ vortex-python/test/test_cuda.py | 24 +++++++ 24 files changed, 353 insertions(+), 63 deletions(-) create mode 100644 vortex-python-cuda/Cargo.toml create mode 100644 vortex-python-cuda/pyproject.toml create mode 100644 vortex-python-cuda/python/vortex_cuda/__init__.py create mode 100644 vortex-python-cuda/python/vortex_cuda/_lib.pyi create mode 100644 vortex-python-cuda/python/vortex_cuda/py.typed create mode 100644 vortex-python-cuda/src/lib.rs create mode 100644 vortex-python-cuda/test/test_cuda.py create mode 100644 vortex-python/test/test_cuda.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a0b516a8818..e4015ddcb29 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -154,9 +154,9 @@ jobs: run: | # Build docs for the whole workspace (aside from vortex-python that is having hard time with the python-docs syntax), # including all private docs. - cargo doc --profile ci --no-deps --document-private-items --workspace --exclude vortex-python + cargo doc --profile ci --no-deps --document-private-items --workspace --exclude vortex-python --exclude vortex-python-cuda # nextest doesn't support doc tests, so we run it here - cargo test --profile ci --doc --workspace --all-features --exclude vortex-cxx --exclude vortex-jni --exclude vortex-ffi --exclude xtask --no-fail-fast + cargo test --profile ci --doc --workspace --all-features --exclude vortex-cxx --exclude vortex-jni --exclude vortex-ffi --exclude xtask --exclude vortex-python-cuda --no-fail-fast build-rust: name: "Rust build (${{matrix.config.name}})" @@ -187,7 +187,7 @@ jobs: target: wasm32-unknown-unknown env: rustflags: "RUSTFLAGS='-A warnings --cfg getrandom_backend=\"unsupported\"'" - args: "--target wasm32-unknown-unknown --exclude vortex --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp --exclude vortex-datafusion --exclude vortex-duckdb --exclude vortex-tui --exclude vortex-zstd --exclude vortex-test-e2e-cuda --exclude vortex-sqllogictest --exclude vortex-parquet-variant" + args: "--target wasm32-unknown-unknown --exclude vortex --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp --exclude vortex-datafusion --exclude vortex-duckdb --exclude vortex-tui --exclude vortex-zstd --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda --exclude vortex-sqllogictest --exclude vortex-parquet-variant" steps: - uses: runs-on/action@v2 if: github.repository == 'vortex-data/vortex' @@ -355,7 +355,7 @@ jobs: --exclude vortex-bench ` --exclude vortex-python --exclude vortex-duckdb ` --exclude vortex-fuzz --exclude vortex-cuda --exclude vortex-cuda-ffi ` - --exclude vortex-nvcomp --exclude vortex-cub --exclude vortex-test-e2e-cuda ` + --exclude vortex-nvcomp --exclude vortex-cub --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda ` --exclude duckdb-bench ` --exclude lance-bench --exclude datafusion-bench --exclude random-access-bench ` --exclude compress-bench --exclude xtask --exclude vortex-datafusion ` diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index d75c711520a..2cdcbaa9dbf 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -100,10 +100,6 @@ jobs: --bench fsst_cuda \ --bench runend_cuda \ --profile bench - - name: Verify CUDA kernels were built - run: test -f vortex-cuda/kernels/gen/release/dynamic_dispatch.ptx - - name: Stage CUDA kernels with Codspeed artifacts - run: cp -R vortex-cuda/kernels/gen/release target/codspeed/cuda-kernels - name: Upload benchmark executables uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 with: @@ -156,7 +152,6 @@ jobs: uses: CodSpeedHQ/action@d872884a306dd4853acf0f584f4b706cf0cc72a2 env: CARGO_MANIFEST_DIR: ${{ github.workspace }}/vortex-cuda - VORTEX_CUDA_KERNELS_DIR: ${{ github.workspace }}/target/codspeed/cuda-kernels with: run: cargo codspeed run $(printf -- '--bench %s ' ${{ matrix.benches }}) token: ${{ secrets.CODSPEED_TOKEN }} diff --git a/.github/workflows/cuda.yaml b/.github/workflows/cuda.yaml index 6ecefc43d43..2bd9afbbdf1 100644 --- a/.github/workflows/cuda.yaml +++ b/.github/workflows/cuda.yaml @@ -38,7 +38,13 @@ jobs: filters: | cuda: - "vortex-cuda/**" + - "vortex-python/python/vortex/__init__.py" + - "vortex-python/pyproject.toml" + - "vortex-python/test/test_cuda.py" + - "vortex-python-cuda/**" - "vortex-test/**" + - "pyproject.toml" + - "uv.lock" - ".github/workflows/**" cuda-build-lint: @@ -62,7 +68,7 @@ jobs: command: >- cargo build --profile ci --locked --all-features --all-targets -p vortex-cuda -p vortex-cuda-ffi -p vortex-cub -p vortex-nvcomp - -p gpu-scan-cli -p vortex-test-e2e-cuda + -p gpu-scan-cli -p vortex-test-e2e-cuda -p vortex-python-cuda - name: Clippy CUDA crates run: | cargo clippy --profile ci --locked --all-features --all-targets \ @@ -72,6 +78,7 @@ jobs: -p vortex-nvcomp \ -p gpu-scan-cli \ -p vortex-test-e2e-cuda \ + -p vortex-python-cuda \ -- -D warnings cuda-test: @@ -95,6 +102,10 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + - name: Install uv + uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + with: + sync: false - name: Install nextest uses: taiki-e/install-action@9e1e5806d4a4822de933115878265be9aaa786d9 # v2 with: @@ -112,10 +123,19 @@ jobs: -p vortex-cub \ -p vortex-nvcomp \ -p vortex-test-e2e-cuda \ + -p vortex-python-cuda \ --all-features \ --no-fail-fast \ --target x86_64-unknown-linux-gnu \ --verbose + - name: Python CUDA extension checks + env: + MATURIN_PEP517_ARGS: "--profile ci" + run: | + # --all-packages installs the shared dev tooling (basedpyright) from the + # root `dev` group; --extra cuda adds the vortex-data-cuda extension. + uv run --all-packages --extra cuda basedpyright vortex-python vortex-python-cuda + uv run --all-packages --extra cuda pytest --benchmark-disable vortex-python/test/test_cuda.py vortex-python-cuda/test cuda-test-sanitizer: needs: [changes] diff --git a/.github/workflows/musl.yml b/.github/workflows/musl.yml index 1078ffe7b0e..6ab9ef566f4 100644 --- a/.github/workflows/musl.yml +++ b/.github/workflows/musl.yml @@ -66,7 +66,7 @@ jobs: run: | cargo nextest run --cargo-profile ci --locked --workspace --no-fail-fast \ --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp \ - --exclude gpu-scan-cli --exclude vortex-test-e2e-cuda \ + --exclude gpu-scan-cli --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda \ --exclude vortex-duckdb --exclude duckdb-bench --exclude vortex-sqllogictest \ --exclude vortex-bench --exclude lance-bench --exclude datafusion-bench --exclude vortex-datafusion \ --exclude compress-bench --exclude random-access-bench --exclude vortex-bench-server diff --git a/.github/workflows/package.yml b/.github/workflows/package.yml index 625bc72d12d..912a36efeb9 100644 --- a/.github/workflows/package.yml +++ b/.github/workflows/package.yml @@ -50,6 +50,30 @@ jobs: run: | cargo set-version --workspace ${{ inputs.version }} + # cargo set-version bumps Cargo.toml crate versions but not the ==0.1.0 pins + # between these two wheels in their pyproject.toml. The base and CUDA extension + # wheels must require each other at the exact released version (they exchange + # Vortex IPC buffers), so rewrite those pins here. + - name: Set Python extension version pins + env: + VERSION: ${{ inputs.version }} + run: | + python3 - <<'PY' + import os + import re + from pathlib import Path + + version = os.environ["VERSION"] + replacements = { + Path("vortex-python/pyproject.toml"): "vortex-data-cuda", + Path("vortex-python-cuda/pyproject.toml"): "vortex-data", + } + for path, package in replacements.items(): + text = path.read_text() + text = re.sub(f"{package}==[^\"\\]]+", f"{package}=={version}", text) + path.write_text(text) + PY + - name: Set up Python uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6 # Latest macOS doesn't allow maturin to install stuff into the global Python interpreter diff --git a/.github/workflows/rust-instrumented.yml b/.github/workflows/rust-instrumented.yml index 57d21964dd6..bf45fc7be13 100644 --- a/.github/workflows/rust-instrumented.yml +++ b/.github/workflows/rust-instrumented.yml @@ -88,7 +88,7 @@ jobs: - name: Rust Tests if: ${{ matrix.suite == 'tests' }} run: | - cargo nextest run --locked --workspace --all-features --no-fail-fast + cargo nextest run --locked --workspace --all-features --no-fail-fast --exclude vortex-python-cuda - name: Generate coverage report run: | RUST_SYSROOT="$(rustc --print sysroot)" diff --git a/Cargo.lock b/Cargo.lock index c7c50f5047e..46c7b93352b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10336,6 +10336,14 @@ dependencies = [ "vortex-tui", ] +[[package]] +name = "vortex-python-cuda" +version = "0.1.0" +dependencies = [ + "pyo3", + "vortex-cuda", +] + [[package]] name = "vortex-row" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 0a1cd5731e9..deed8b8d58a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ members = [ "fuzz", "vortex-jni", "vortex-python", + "vortex-python-cuda", "vortex-tui", "vortex-web/crate", "vortex-sqllogictest", diff --git a/docs/api/python/runtime.rst b/docs/api/python/runtime.rst index 1c404c74d32..b5307a4041b 100644 --- a/docs/api/python/runtime.rst +++ b/docs/api/python/runtime.rst @@ -9,6 +9,7 @@ otherwise to the number of available CPU cores minus one. Use .. autosummary:: :nosignatures: + ~vortex.cuda_extension_installed ~vortex.set_worker_threads ~vortex.worker_threads @@ -16,6 +17,8 @@ otherwise to the number of available CPU cores minus one. Use
+.. autofunction:: vortex.cuda_extension_installed + .. autofunction:: vortex.set_worker_threads .. autofunction:: vortex.worker_threads diff --git a/pyproject.toml b/pyproject.toml index f0e82ec52c9..53b769f350c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,11 @@ members = ["bench-orchestrator", "vortex-python", "docs"] [tool.uv.sources] bench-orchestrator = { workspace = true } vortex-data = { workspace = true } +# Resolve vortex-data-cuda from its local path so the optional `cuda` extra is +# lockable before the package is published to PyPI. A path source (not +# `workspace = true`) keeps the CUDA extension out of the default `--all-packages` +# set, so it is only built when the `cuda` extra is requested. +vortex-data-cuda = { path = "vortex-python-cuda" } docs = { workspace = true } [tool.ruff] diff --git a/uv.lock b/uv.lock index bfc71cb994e..8748a91544e 100644 --- a/uv.lock +++ b/uv.lock @@ -1899,6 +1899,9 @@ dependencies = [ ] [package.optional-dependencies] +cuda = [ + { name = "vortex-data-cuda" }, +] duckdb = [ { name = "duckdb" }, ] @@ -1941,8 +1944,9 @@ requires-dist = [ { name = "ray", marker = "extra == 'ray'", specifier = ">=2.48" }, { name = "substrait", specifier = ">=0.23.0,<0.29.0" }, { name = "typing-extensions", specifier = ">=4.5.0" }, + { name = "vortex-data-cuda", marker = "extra == 'cuda'", directory = "vortex-python-cuda" }, ] -provides-extras = ["duckdb", "numpy", "pandas", "polars", "ray"] +provides-extras = ["cuda", "duckdb", "numpy", "pandas", "polars", "ray"] [package.metadata.requires-dev] dev = [ @@ -1960,6 +1964,16 @@ dev = [ { name = "ray", specifier = ">=2.48" }, ] +[[package]] +name = "vortex-data-cuda" +source = { directory = "vortex-python-cuda" } +dependencies = [ + { name = "vortex-data" }, +] + +[package.metadata] +requires-dist = [{ name = "vortex-data", specifier = "==0.1.0" }] + [[package]] name = "watchfiles" version = "1.1.0" diff --git a/vortex-cuda/build.rs b/vortex-cuda/build.rs index eb4743bab21..e201345cd73 100644 --- a/vortex-cuda/build.rs +++ b/vortex-cuda/build.rs @@ -8,6 +8,7 @@ use std::env; use std::fs::File; use std::io; +use std::io::Write; use std::path::Path; use std::path::PathBuf; use std::process::Command; @@ -32,15 +33,6 @@ fn main() { std::fs::create_dir_all(&kernels_gen).expect("Failed to create kernels/gen directory"); - // Always emit the kernels output directory path as a compile-time env var so any binary - // linking against vortex-cuda can find the PTX files. This must be set regardless - // of CUDA availability since the code using env!() is always compiled. - // At runtime, VORTEX_CUDA_KERNELS_DIR can be set to override this path. - println!( - "cargo:rustc-env=VORTEX_CUDA_KERNELS_DIR={}", - kernels_gen.display() - ); - println!("cargo:rerun-if-env-changed=PROFILE"); // Regenerate bit_unpack kernels only when the generator changes @@ -60,6 +52,8 @@ fn main() { generate_dynamic_dispatch_bindings(&kernels_src, &out_dir); generate_patches_bindings(&kernels_src, &out_dir); + generate_embedded_ptx(&out_dir, &kernels_gen).expect("Failed to generate embedded PTX source"); + if !is_cuda_available() { return; } @@ -97,6 +91,57 @@ fn main() { } } } + + // Refresh embedded_ptx.rs after nvcc has compiled PTX so the binary embeds the latest kernels. + generate_embedded_ptx(&out_dir, &kernels_gen).expect("Failed to generate embedded PTX source"); +} + +/// Generates `embedded_ptx.rs`, a module-name lookup table of `include_str!`-embedded PTX. +/// +/// Called before the `nvcc` check to create an empty table for no-toolkit builds, then again after +/// PTX compilation so runtime loading never depends on build-machine filesystem paths. +fn generate_embedded_ptx(out_dir: &Path, kernels_gen: &Path) -> io::Result<()> { + let mut ptx_files = Vec::new(); + if let Ok(entries) = std::fs::read_dir(kernels_gen) { + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().and_then(|e| e.to_str()) == Some("ptx") { + ptx_files.push(path); + } + } + } + ptx_files.sort(); + + let mut file = File::create(out_dir.join("embedded_ptx.rs"))?; + if ptx_files.is_empty() { + writeln!( + file, + "pub(crate) fn embedded_ptx(_module_name: &str) -> Option<&'static str> {{" + )?; + writeln!(file, " None")?; + writeln!(file, "}}")?; + return Ok(()); + } + + writeln!( + file, + "pub(crate) fn embedded_ptx(module_name: &str) -> Option<&'static str> {{" + )?; + writeln!(file, " match module_name {{")?; + for path in ptx_files { + let Some(module_name) = path.file_stem().and_then(|stem| stem.to_str()) else { + continue; + }; + writeln!( + file, + " {module_name:?} => Some(include_str!({:?})),", + path.to_string_lossy() + )?; + } + writeln!(file, " _ => None,")?; + writeln!(file, " }}")?; + writeln!(file, "}}")?; + Ok(()) } fn generate_unpack(output_dir: &Path, thread_count: usize) -> io::Result { diff --git a/vortex-cuda/src/kernel/mod.rs b/vortex-cuda/src/kernel/mod.rs index 5665bc6c3f3..323458db2d8 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -4,8 +4,6 @@ //! CUDA kernel loading and management. use std::fmt::Debug; -use std::path::Path; -use std::path::PathBuf; use std::sync::Arc; use cudarc::driver::CudaContext; @@ -22,6 +20,10 @@ use vortex::error::vortex_err; use vortex::utils::aliases::dash_map::DashMap; mod arrays; +// PTX kernels embedded at build time; generated by build.rs's generate_embedded_ptx. +mod embedded_ptx { + include!(concat!(env!("OUT_DIR"), "/embedded_ptx.rs")); +} mod encodings; mod filter; mod patches; @@ -235,24 +237,9 @@ impl KernelLoader { let module = if let Some(entry) = self.modules.get(module_name) { Arc::clone(entry.value()) } else { - let ptx_path = Self::ptx_path_for_module(module_name); - - // Compile and load the CUDA module. - let module = cuda_context - .load_module(Ptx::from_file(&ptx_path)) - .map_err(|e| { - vortex_err!( - "Failed to load CUDA module {}, ptx path {}: {}", - module_name, - ptx_path.display(), - e - ) - })?; - - // Cache the module + let module = Self::load_module(module_name, cuda_context)?; self.modules .insert(module_name.to_string(), Arc::clone(&module)); - module }; @@ -262,22 +249,18 @@ impl KernelLoader { .map_err(|e| vortex_err!("Failed to load kernel function '{}': {}", kernel_name, e)) } - /// Returns the PTX file path for a given module name. - /// - /// Checks for `VORTEX_CUDA_KERNELS_DIR` environment variable at runtime first, - /// falling back to the path baked in at compile time by build.rs. - /// - /// # Arguments - /// - /// * `module_name` - Name of the module - /// - /// # Returns - /// - /// The full path to the PTX file - fn ptx_path_for_module(module_name: &str) -> PathBuf { - let kernels_dir = std::env::var("VORTEX_CUDA_KERNELS_DIR") - .unwrap_or_else(|_| env!("VORTEX_CUDA_KERNELS_DIR").to_string()); - Path::new(&kernels_dir).join(format!("{}.ptx", module_name)) + /// Loads a CUDA module from PTX embedded in the binary. + fn load_module( + module_name: &str, + cuda_context: &Arc, + ) -> VortexResult> { + let ptx = embedded_ptx::embedded_ptx(module_name).ok_or_else(|| { + vortex_err!("CUDA module {module_name} was not embedded at build time") + })?; + + cuda_context + .load_module(Ptx::from_src(ptx)) + .map_err(|e| vortex_err!("Failed to load embedded CUDA module {module_name}: {e}")) } } diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 021c811a90e..e5c00d0b3d0 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -3,8 +3,8 @@ //! CUDA support for Vortex arrays. -use std::process::Command; - +use cudarc::driver::CudaContext; +use cudarc::driver::sys; use tracing::info; pub mod arrow; @@ -89,12 +89,28 @@ pub use vortex_nvcomp as nvcomp; use crate::kernel::SequenceExecutor; use crate::kernel::SliceExecutor; -/// Checks if CUDA is available on the system by looking for nvcc. +#[cfg(test)] +pub(crate) fn canonicalize_cpu( + array: impl vortex::array::IntoArray, +) -> vortex::error::VortexResult { + use vortex::array::LEGACY_SESSION; + use vortex::array::VortexSessionExecute; + + array + .into_array() + .execute::(&mut LEGACY_SESSION.create_execution_ctx()) +} + +/// Checks if a CUDA driver and at least one CUDA device are available. +/// +/// cudarc loads `libcuda` lazily and panics if the driver library is absent, so we first probe +/// for it with cudarc's own `is_culib_present`. Creating the context then fails gracefully with +/// `Err`, rather than panicking, when the driver is present but no usable device is. pub fn cuda_available() -> bool { - Command::new("nvcc") - .arg("--version") - .output() - .is_ok_and(|o| o.status.success()) + // SAFETY: `is_culib_present` only tries to dlopen the driver library to test for its + // presence; it upholds no further invariants. + let driver_present = unsafe { sys::is_culib_present() }; + driver_present && CudaContext::new(0).is_ok() } /// Registers CUDA kernels. diff --git a/vortex-python-cuda/Cargo.toml b/vortex-python-cuda/Cargo.toml new file mode 100644 index 00000000000..8735175177d --- /dev/null +++ b/vortex-python-cuda/Cargo.toml @@ -0,0 +1,31 @@ +[package] +name = "vortex-python-cuda" +authors = { workspace = true } +categories = { workspace = true } +description = "Optional CUDA extension for PyVortex." +edition = { workspace = true } +homepage = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +publish = false +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[lints] +workspace = true + +[lib] +crate-type = ["cdylib"] + +[features] +default = ["extension-module"] +# Maturin enables `pyo3/extension-module` via pyproject.toml when building wheels. Keep +# this local feature empty so plain `cargo build`/`cargo check` links libpython normally. +extension-module = [] + +[dependencies] +pyo3 = { workspace = true, features = ["abi3", "abi3-py311"] } +vortex-cuda = { workspace = true } diff --git a/vortex-python-cuda/pyproject.toml b/vortex-python-cuda/pyproject.toml new file mode 100644 index 00000000000..5ee577e4875 --- /dev/null +++ b/vortex-python-cuda/pyproject.toml @@ -0,0 +1,34 @@ +[project] +name = "vortex-data-cuda" +# Maturin copies these, at build time, from Cargo.toml. +dynamic = ["version", "description", "authors"] +requires-python = ">= 3.11" +# The CUDA extension package must exactly match the base package version because the two extensions +# exchange Vortex IPC buffers. Keep this in sync with the workspace package version; release +# automation may rewrite it when publishing a non-workspace version. +dependencies = ["vortex-data==0.1.0"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", + "Topic :: Database", + "Topic :: Scientific/Engineering", + "Environment :: GPU :: NVIDIA CUDA", +] + +[build-system] +requires = ["maturin>=1.10.2,<2.0"] +build-backend = "maturin" + +[tool.uv] +managed = true + +[tool.maturin] +profile = "release" +editable-profile = "dev" +python-source = "python" +module-name = "vortex_cuda._lib" +features = ["pyo3/extension-module"] +compatibility = "manylinux2014" diff --git a/vortex-python-cuda/python/vortex_cuda/__init__.py b/vortex-python-cuda/python/vortex_cuda/__init__.py new file mode 100644 index 00000000000..d259d24d698 --- /dev/null +++ b/vortex-python-cuda/python/vortex_cuda/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +from ._lib import cuda_available # pyright: ignore[reportMissingModuleSource] + +__all__ = ["cuda_available"] diff --git a/vortex-python-cuda/python/vortex_cuda/_lib.pyi b/vortex-python-cuda/python/vortex_cuda/_lib.pyi new file mode 100644 index 00000000000..431190faa0a --- /dev/null +++ b/vortex-python-cuda/python/vortex_cuda/_lib.pyi @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +def cuda_available() -> bool: ... diff --git a/vortex-python-cuda/python/vortex_cuda/py.typed b/vortex-python-cuda/python/vortex_cuda/py.typed new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vortex-python-cuda/src/lib.rs b/vortex-python-cuda/src/lib.rs new file mode 100644 index 00000000000..26da87f2530 --- /dev/null +++ b/vortex-python-cuda/src/lib.rs @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Optional CUDA extension for PyVortex. +//! +//! Builds the separate `vortex-data-cuda` wheel (imported as `vortex_cuda`), installed alongside +//! the CPU-only `vortex-data` wheel. Keeping CUDA in its own extension keeps the base wheel free of +//! CUDA build/runtime dependencies; `vortex.cuda_extension_installed()` reports whether it is present. + +use pyo3::prelude::*; + +/// Return whether a usable CUDA device is available in the current process. +/// +/// This performs a runtime probe of the CUDA driver and device. It differs from +/// `vortex.cuda_extension_installed()`, which only reports whether this extension package is +/// installed. +#[pyfunction] +fn cuda_available() -> bool { + vortex_cuda::cuda_available() +} + +/// The `vortex_cuda._lib` extension module. +#[cfg(feature = "extension-module")] +#[pymodule] +fn _lib(m: &Bound) -> PyResult<()> { + m.add_function(wrap_pyfunction!(cuda_available, m)?)?; + Ok(()) +} diff --git a/vortex-python-cuda/test/test_cuda.py b/vortex-python-cuda/test/test_cuda.py new file mode 100644 index 00000000000..de269164317 --- /dev/null +++ b/vortex-python-cuda/test/test_cuda.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +import tomllib +from pathlib import Path +from typing import cast + +import vortex_cuda + +import vortex + + +def workspace_version() -> str: + workspace_pyproject = tomllib.loads((Path(__file__).parents[2] / "Cargo.toml").read_text()) + return cast(str, workspace_pyproject["workspace"]["package"]["version"]) + + +def test_extension_is_detected_by_base(): + assert vortex.cuda_extension_installed() is True + + +def test_cuda_available_returns_bool(): + assert isinstance(vortex_cuda.cuda_available(), bool) + + +def test_extension_exact_pins_base_package(): + pyproject = tomllib.loads((Path(__file__).parents[1] / "pyproject.toml").read_text()) + + assert pyproject["project"]["dependencies"] == [f"vortex-data=={workspace_version()}"] diff --git a/vortex-python/pyproject.toml b/vortex-python/pyproject.toml index ef8e76373a6..8ee24e94fa5 100644 --- a/vortex-python/pyproject.toml +++ b/vortex-python/pyproject.toml @@ -37,6 +37,7 @@ classifiers = [ vx = "vortex.cli:main" [project.optional-dependencies] +cuda = ["vortex-data-cuda==0.1.0"] polars = ["polars>=1.31.0"] pandas = ["pandas>=2.2.0"] numpy = ["numpy>=1.26.0"] diff --git a/vortex-python/python/vortex/__init__.py b/vortex-python/python/vortex/__init__.py index 9cea6285f2b..ce4e57064ff 100644 --- a/vortex-python/python/vortex/__init__.py +++ b/vortex-python/python/vortex/__init__.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright the Vortex contributors import importlib.metadata +import importlib.util from . import _lib, arrays, dataset, expr, file, io, ray, registry, scan from ._lib.arrays import ( # pyright: ignore[reportMissingModuleSource] @@ -100,6 +101,23 @@ # If the distribution is not installed, keep the unknown fallback. pass + +def cuda_extension_installed() -> bool: + """Return whether the optional Vortex CUDA extension package is importable. + + The base ``vortex-data`` wheel is CPU-only. Optional CUDA functionality is + provided by the separate ``vortex-data-cuda`` extension package. This returns + ``True`` when the ``vortex_cuda`` import package can be found in the current + environment, which is what ``vortex-data[cuda]`` installs. + + This does not probe the CUDA driver or attached devices, and it does not + imply that any particular GPU interop API is available. After installing the + extension package, use ``vortex_cuda.cuda_available()`` to check whether CUDA + is usable at runtime. + """ + return importlib.util.find_spec("vortex_cuda") is not None + + __all__ = [ # --- Modules --- "arrays", @@ -113,6 +131,7 @@ # --- Objects and Functions --- "array", "compress", + "cuda_extension_installed", # Arrays "Array", "PyArray", diff --git a/vortex-python/test/test_cuda.py b/vortex-python/test/test_cuda.py new file mode 100644 index 00000000000..5c6bed3c645 --- /dev/null +++ b/vortex-python/test/test_cuda.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +import tomllib +from pathlib import Path +from typing import cast + +import vortex + + +def workspace_version() -> str: + workspace_pyproject = tomllib.loads((Path(__file__).parents[2] / "Cargo.toml").read_text()) + return cast(str, workspace_pyproject["workspace"]["package"]["version"]) + + +def test_cuda_extension_installed_returns_bool(): + assert isinstance(vortex.cuda_extension_installed(), bool) + assert "cuda_extension_installed" in vortex.__all__ + + +def test_cuda_extra_installs_exact_matching_extension(): + pyproject = tomllib.loads((Path(__file__).parents[1] / "pyproject.toml").read_text()) + + assert pyproject["project"]["optional-dependencies"]["cuda"] == [f"vortex-data-cuda=={workspace_version()}"] From 4c7eb01abba1a7d73c9927684a2ffb712880e62c Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Wed, 24 Jun 2026 14:41:22 +0200 Subject: [PATCH 2/3] fail explicitly in case the version parsing fails Signed-off-by: Alexander Droste --- .github/workflows/package.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/package.yml b/.github/workflows/package.yml index 912a36efeb9..3de9ac5c81d 100644 --- a/.github/workflows/package.yml +++ b/.github/workflows/package.yml @@ -70,7 +70,11 @@ jobs: } for path, package in replacements.items(): text = path.read_text() - text = re.sub(f"{package}==[^\"\\]]+", f"{package}=={version}", text) + text, count = re.subn( + f"{re.escape(package)}==[^\"\\]]+", f"{package}=={version}", text + ) + if count != 1: + raise RuntimeError(f"Expected one {package} version pin in {path}, found {count}") path.write_text(text) PY From c8ecde6fefa44397056651bb694fecb74583cca0 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Wed, 24 Jun 2026 15:05:06 +0200 Subject: [PATCH 3/3] drop Signed-off-by: Alexander Droste --- vortex-cuda/src/lib.rs | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index e5c00d0b3d0..07cfe40b06f 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -89,18 +89,6 @@ pub use vortex_nvcomp as nvcomp; use crate::kernel::SequenceExecutor; use crate::kernel::SliceExecutor; -#[cfg(test)] -pub(crate) fn canonicalize_cpu( - array: impl vortex::array::IntoArray, -) -> vortex::error::VortexResult { - use vortex::array::LEGACY_SESSION; - use vortex::array::VortexSessionExecute; - - array - .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) -} - /// Checks if a CUDA driver and at least one CUDA device are available. /// /// cudarc loads `libcuda` lazily and panics if the driver library is absent, so we first probe