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..3de9ac5c81d 100644
--- a/.github/workflows/package.yml
+++ b/.github/workflows/package.yml
@@ -50,6 +50,34 @@ 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, 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
+
- 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..07cfe40b06f 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,16 @@ 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.
+/// 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()}"]