feat(python): package CUDA as an optional extension (#8510)

Keeps 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 extension is importable, while
`vortex_cuda.cuda_available()` performs the runtime CUDA probe.

Wires `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.

Keeps `vortex-python-cuda` out of the uv workspace so `uv
--all-packages` on non-GPU CI does not build the CUDA crate chain. Add
explicit CUDA Python checks to the GPU workflow instead, including the
uv setup needed to run them.

Embeds generated PTX into `vortex-cuda` when available so installed
wheels do not depend on build-machine kernel file paths at runtime.

---------

Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>

Author: 0ax1

.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 `

.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 }}

.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]

.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

.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

.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)"

Cargo.lock

@@ -37,6 +37,7 @@ members = [
     "fuzz",
     "vortex-jni",
     "vortex-python",
+    "vortex-python-cuda",
     "vortex-tui",
     "vortex-web/crate",
     "vortex-sqllogictest",

Cargo.toml

@@ -9,13 +9,16 @@ otherwise to the number of available CPU cores minus one. Use
 .. autosummary::
    :nosignatures:
 
+   ~vortex.cuda_extension_installed
    ~vortex.set_worker_threads
    ~vortex.worker_threads
 
 .. raw:: html
 
    <hr>
 
+.. autofunction:: vortex.cuda_extension_installed
+
 .. autofunction:: vortex.set_worker_threads
 
 .. autofunction:: vortex.worker_threads

docs/api/python/runtime.rst

@@ -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]