diff --git a/.github/workflows/bench-dispatch.yml b/.github/workflows/bench-dispatch.yml index d1116d964b5..2e704a080f1 100644 --- a/.github/workflows/bench-dispatch.yml +++ b/.github/workflows/bench-dispatch.yml @@ -32,6 +32,22 @@ jobs: uses: ./.github/workflows/bench-pr.yml secrets: inherit + remove-gpu-compress-bench-label: + runs-on: ubuntu-latest + timeout-minutes: 10 + if: github.event.label.name == 'action/benchmark-gpu-compress' + steps: + - uses: actions-ecosystem/action-remove-labels@2ce5d41b4b6aa8503e285553f75ed56e0a40bae0 # v1 + if: github.event.pull_request.head.repo.full_name == 'vortex-data/vortex' + with: + labels: action/benchmark-gpu-compress + fail_on_error: true + + gpu-compress-bench: + needs: remove-gpu-compress-bench-label + uses: ./.github/workflows/gpu-compress-bench-pr.yml + secrets: inherit + remove-sql-label: runs-on: ubuntu-latest timeout-minutes: 10 diff --git a/.github/workflows/gpu-compress-bench-pr.yml b/.github/workflows/gpu-compress-bench-pr.yml new file mode 100644 index 00000000000..58654974fdd --- /dev/null +++ b/.github/workflows/gpu-compress-bench-pr.yml @@ -0,0 +1,87 @@ +# Runs the GPU-enabled compression decompression benchmarks for a pull request. +# Called from bench-dispatch.yml when the `action/benchmark-gpu-compress` label is added. + +name: GPU Compression Benchmarks + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: false + +on: + workflow_call: { } + workflow_dispatch: { } + +permissions: + contents: read + pull-requests: write + +jobs: + gpu-compress-bench: + name: GPU Compression + timeout-minutes: 120 + runs-on: >- + runs-on=${{ github.run_id }}/family=g5/cpu=8/image=ubuntu24-gpu-x64/extras=s3-cache/tag=gpu-compress-bench + steps: + - uses: runs-on/action@v2 + with: + sccache: s3 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + with: + ref: ${{ github.event.pull_request.head.sha || github.sha }} + - name: Setup benchmark environment + run: sudo bash scripts/setup-benchmark.sh + - uses: ./.github/actions/setup-rust + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: "true" + - uses: ./.github/actions/system-info + - name: Display NVIDIA GPU details + run: | + nvidia-smi + nvidia-smi -L + nvidia-smi -q -d Memory + - name: Build GPU compression benchmark + shell: bash + run: | + cargo build --locked --package compress-bench --profile release_debug --features cuda,unstable_encodings + - name: Run GPU compression benchmark + shell: bash + env: + RUST_BACKTRACE: full + # Do not enable VORTEX_EXPERIMENTAL_PATCHED_ARRAY here: it rewrites interior + # `Patches` into a top-level `vortex.patched` array, which has no CUDA decode + # kernel and makes GPU decompression fail with "No CUDA kernel for encoding + # vortex.patched". Leaving it unset keeps patches inline in the bitpacked array, + # which the CUDA bitpacked kernel decodes. This matches the GPU test workflow + # (cuda.yaml), which also sets only FLAT_LAYOUT_INLINE_ARRAY_NODE. + FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + run: | + bash scripts/bench-taskset.sh target/release_debug/compress-bench \ + --gpu-decompress -d table > gpu-compress.txt + cat gpu-compress.txt + - name: Publish results + shell: bash + run: | + { + echo "# GPU Compression" + echo + echo '```text' + cat gpu-compress.txt + echo '```' + } > comment.md + cat comment.md >> "$GITHUB_STEP_SUMMARY" + - name: Comment PR + if: github.event.pull_request.head.repo.fork == false + uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b # v3 + with: + file-path: comment.md + comment-tag: bench-pr-comment-gpu-compress + - name: Comment PR on failure + if: failure() && github.event.pull_request.head.repo.fork == false + uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b # v3 + with: + message: | + # BENCHMARK FAILED + + GPU Compression failed. Check the [workflow run](${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}) for details. + comment-tag: bench-pr-comment-gpu-compress diff --git a/Cargo.lock b/Cargo.lock index b562097d72a..f4ea25fc531 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1556,11 +1556,13 @@ dependencies = [ "lance-bench", "parquet 58.3.0", "regex", + "tempfile", "tokio", "tracing", "vortex", "vortex-arrow", "vortex-bench", + "vortex-cuda", ] [[package]] @@ -9782,6 +9784,7 @@ dependencies = [ "rand 0.10.2", "regex", "reqwest 0.13.4", + "rstest", "serde", "serde_json", "sha2 0.11.0", diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index 4e8fbc530b2..4046a12d42e 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -27,15 +27,18 @@ itertools = { workspace = true } lance-bench = { path = "../lance-bench", optional = true } parquet = { workspace = true } regex = { workspace = true } +tempfile = { workspace = true, optional = true } tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } vortex = { workspace = true } vortex-arrow = { workspace = true } vortex-bench = { workspace = true } +vortex-cuda = { workspace = true, optional = true } [features] +cuda = ["dep:tempfile", "dep:vortex-cuda"] lance = ["dep:lance-bench"] -unstable_encodings = ["vortex/unstable_encodings"] +unstable_encodings = ["vortex/unstable_encodings", "vortex-cuda?/unstable_encodings"] [[bin]] name = "compress-bench" diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index df1027e2bd0..bf2d3efc1db 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -14,3 +14,14 @@ See [`src/main.rs`](./src/main.rs) for the dataset list and CLI flags (`--format ```bash cargo run -p compress-bench --profile release_debug ``` + +GPU decompression is opt-in and runs only the existing benchmark names allow-listed in +`src/main.rs`: + +```bash +cargo run -p compress-bench --profile release_debug \ + --features cuda,unstable_encodings -- --gpu-decompress +``` + +On Linux, GPU files are read with direct IO (`O_DIRECT`) so repeated iterations measure +storage bandwidth rather than page-cache hits. diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs new file mode 100644 index 00000000000..3dbb68bc7a8 --- /dev/null +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::hint::black_box; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Result; +use async_trait::async_trait; +use futures::StreamExt; +use tempfile::NamedTempFile; +use vortex::array::IntoArray; +use vortex::array::arrays::StructArray; +use vortex::array::arrays::struct_::StructArrayExt; +use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex_bench::Format; +use vortex_bench::SESSION; +use vortex_bench::compress::Compressor; +use vortex_bench::conversions::parquet_to_vortex_chunks; +use vortex_cuda::CudaOpenOptionsExt; +use vortex_cuda::CudaSession; +#[cfg(target_os = "linux")] +use vortex_cuda::PooledFileReadAtOptions; +use vortex_cuda::executor::CudaArrayExt; +use vortex_cuda::layout::CudaFlatLayoutStrategy; +use vortex_cuda::layout::register_cuda_layout; + +/// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU. +pub struct GpuVortexCompressor; + +#[async_trait] +impl Compressor for GpuVortexCompressor { + fn format(&self) -> Format { + Format::OnDiskVortex + } + + async fn compress(&self, _parquet_path: &Path) -> Result<(u64, Duration)> { + anyhow::bail!("GPU compress-bench only supports decompression measurements") + } + + async fn decompress(&self, parquet_path: &Path) -> Result { + register_cuda_layout(&SESSION); + + let uncompressed = parquet_to_vortex_chunks(parquet_path.to_path_buf()).await?; + let gpu_file = NamedTempFile::new()?; + let mut output = tokio::fs::File::create(gpu_file.path()).await?; + let strategy = WriteStrategyBuilder::default() + .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().only_cuda_compatible()) + .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) + .build(); + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut output, uncompressed.into_array().to_array_stream()) + .await?; + output.sync_all().await?; + drop(output); + + let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; + let start = Instant::now(); + let open_options = SESSION.open_options().with_cuda(); + // Direct IO keeps repeated iterations measuring storage bandwidth rather than + // page-cache hits. It is only available on Linux. + #[cfg(target_os = "linux")] + let open_options = + open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io()); + let file = open_options.open_path(gpu_file.path()).await?; + let mut batches = file.scan()?.into_array_stream()?; + + while let Some(batch) = batches.next().await { + let record = batch?.execute::(cuda_ctx.execution_ctx())?; + for field in record.iter_unmasked_fields() { + black_box(field.clone().execute_cuda(&mut cuda_ctx).await?); + } + } + cuda_ctx.synchronize_stream()?; + + Ok(start.elapsed()) + } +} diff --git a/benchmarks/compress-bench/src/lib.rs b/benchmarks/compress-bench/src/lib.rs index b1692c40975..68039996605 100644 --- a/benchmarks/compress-bench/src/lib.rs +++ b/benchmarks/compress-bench/src/lib.rs @@ -3,5 +3,7 @@ #[cfg(feature = "lance")] pub use lance_bench::compress::LanceCompressor; +#[cfg(feature = "cuda")] +pub mod gpu_vortex; pub mod parquet; pub mod vortex; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 1cbf3f47c55..97ead639325 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -7,6 +7,8 @@ use std::time::Duration; use clap::Parser; #[cfg(feature = "lance")] use compress_bench::LanceCompressor; +#[cfg(feature = "cuda")] +use compress_bench::gpu_vortex::GpuVortexCompressor; use compress_bench::parquet::ParquetCompressor; use compress_bench::vortex::VortexCompressor; use indicatif::ProgressBar; @@ -65,6 +67,11 @@ struct Args { ops: Vec, #[arg(long)] datasets: Option, + /// Run GPU decompression for the allow-listed benchmarks. + /// + /// This filters the suite to GPU-supported dataset names and runs only Vortex decompression. + #[arg(long)] + gpu_decompress: bool, #[arg(short, long, default_value_t, value_enum)] display_format: DisplayFormat, #[arg(short, long)] @@ -87,11 +94,22 @@ async fn main() -> anyhow::Result<()> { setup_logging_and_tracing_with_format(args.verbose, args.tracing, args.log_format)?; + if args.gpu_decompress && !cfg!(feature = "cuda") { + anyhow::bail!("--gpu-decompress requires building compress-bench with --features cuda"); + } + + let (formats, ops) = if args.gpu_decompress { + (vec![Format::OnDiskVortex], vec![CompressOp::Decompress]) + } else { + (args.formats, args.ops) + }; + run_compress( args.iterations, args.datasets.map(|d| Regex::new(&d)).transpose()?, - args.formats, - args.ops, + formats, + ops, + args.gpu_decompress, args.display_format, args.output_path, args.gh_json_v3, @@ -100,7 +118,16 @@ async fn main() -> anyhow::Result<()> { } /// Get a compressor for the given format. -fn get_compressor(format: Format) -> Box { +fn get_compressor(format: Format, gpu_decompress: bool) -> Box { + if gpu_decompress { + #[cfg(feature = "cuda")] + { + return Box::new(GpuVortexCompressor); + } + #[cfg(not(feature = "cuda"))] + unreachable!("GPU feature validation happens before selecting compressors"); + } + match format { Format::OnDiskVortex => Box::new(VortexCompressor), Format::Parquet => Box::new(ParquetCompressor::new()), @@ -116,11 +143,16 @@ const BENCHMARK_ID: &str = "compress"; /// Repo-relative path of the suite explainer linked from CI benchmark PR comments. const DOC_PATH: &str = "benchmarks/compress-bench/README.md"; +#[expect( + clippy::too_many_arguments, + reason = "benchmark CLI options are forwarded one-to-one" +)] async fn run_compress( iterations: usize, datasets_filter: Option, formats: Vec, ops: Vec, + gpu_decompress: bool, display_format: DisplayFormat, output_path: Option, gh_json_v3: Option, @@ -147,6 +179,14 @@ async fn run_compress( // ), ]; + // Add an existing benchmark name here only after its CUDA-compatible compression and + // decompression kernels have been verified end to end. + #[expect( + clippy::useless_vec, + reason = "this is an intentionally incremental allow-list of benchmark names" + )] + let gpu_decompress_benchmarks = vec!["TPC-H l_comment canonical"]; + let datasets: Vec<&dyn Dataset> = [ &TaxiData as &dyn Dataset, PBI_DATASETS.get(Arade), @@ -168,6 +208,9 @@ async fn run_compress( .into_iter() .chain(structlistofints.iter().map(|d| d as &dyn Dataset)) .filter(|d| { + if gpu_decompress && !gpu_decompress_benchmarks.contains(&d.name()) { + return false; + } if let Some(filter) = datasets_filter.as_ref() { filter.is_match(d.name()) } else { @@ -184,9 +227,15 @@ async fn run_compress( let mut v3_records: Vec = Vec::new(); for dataset_handle in datasets.into_iter() { - let (m, mut records) = - run_benchmark_for_dataset(&progress, &formats, &ops, iterations, dataset_handle) - .await?; + let (m, mut records) = run_benchmark_for_dataset( + &progress, + &formats, + &ops, + iterations, + dataset_handle, + gpu_decompress, + ) + .await?; measurements.push(m); v3_records.append(&mut records); } @@ -227,6 +276,7 @@ async fn run_benchmark_for_dataset( ops: &[CompressOp], iterations: usize, dataset_handle: &dyn Dataset, + gpu_decompress: bool, ) -> anyhow::Result<(CompressMeasurements, Vec)> { let bench_name = dataset_handle.name(); let (v3_dataset, v3_variant) = dataset_handle.v3_dataset_dims(); @@ -242,7 +292,7 @@ async fn run_benchmark_for_dataset( let mut v3_records: Vec = Vec::new(); for format in formats { - let compressor = get_compressor(*format); + let compressor = get_compressor(*format, gpu_decompress); for op in ops { let time = match op { @@ -293,7 +343,11 @@ async fn run_benchmark_for_dataset( v3_records.push(v3::compression_time_record( &result.timing, v3_dataset, - v3_variant, + if gpu_decompress { + Some("gpu") + } else { + v3_variant + }, CompressOp::Decompress, all_runs_ns, )); diff --git a/docs/developer-guide/benchmarking.md b/docs/developer-guide/benchmarking.md index dfa0fa1cd45..1655d7f8a08 100644 --- a/docs/developer-guide/benchmarking.md +++ b/docs/developer-guide/benchmarking.md @@ -212,6 +212,8 @@ Benchmarks run automatically on all commits to `develop` and can be run on-deman `develop`, with results uploaded for historical tracking. - **PR benchmarks** -- triggered by the `action/benchmark` label. Results are compared against the latest `develop` run and posted as a PR comment. +- **GPU compression benchmarks** -- triggered by the `action/benchmark-gpu-compress` label. Runs + the allow-listed Vortex decompression cases on a GPU runner and posts the timings as a PR comment. - **SQL benchmarks** -- triggered by the `action/benchmark-sql` label. Runs the base SQL matrix, which excludes Appian, TPC-H SF=10 on S3, `vortex-compact`, and `duckdb:duckdb`. - **Full SQL benchmarks** -- triggered by the `action/benchmark-sql-full` label. Runs the full diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index d3a23c976f1..6dbe8ffea98 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -83,6 +83,7 @@ wkb = { workspace = true } [dev-dependencies] insta = { workspace = true } +rstest = { workspace = true } [features] unstable_encodings = ["vortex/unstable_encodings"] diff --git a/vortex-bench/src/display.rs b/vortex-bench/src/display.rs index 343c1a06612..97203abdb4d 100644 --- a/vortex-bench/src/display.rs +++ b/vortex-bench/src/display.rs @@ -30,6 +30,16 @@ pub fn render_table( all_measurements: Vec, targets: &[Target], ) -> anyhow::Result<()> { + // `all_measurements` is empty for decompress-only runs such as `compress-bench + // --gpu-decompress`: every ratio compares vortex against parquet or lance, neither of + // which is benchmarked there, so no ratio is recorded even though a baseline target is + // still passed. `targets` is instead empty when the caller has no baseline format to + // report against. Either way there is no table to draw, and the indexing below would + // panic looking up a baseline that has no measurements. + if all_measurements.is_empty() || targets.is_empty() { + return Ok(()); + } + let mut measurements: HashMap> = HashMap::with_capacity(all_measurements.len().div_ceil(targets.len())); @@ -127,3 +137,30 @@ fn color(baseline: MeasurementValue, value: MeasurementValue) -> Color { Color::BG_BRIGHT_GREEN | Color::FG_BLACK } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + use crate::Engine; + use crate::Format; + use crate::measurements::CompressionTimingMeasurement; + + /// Decompress-only runs (`compress-bench --gpu-decompress`) collect no compression ratios, + /// and callers pass an empty target list when the baseline format is absent. Both used to + /// panic: the first on the missing baseline key, the second on a divide-by-zero capacity. + #[rstest] + #[case::no_measurements(&[Target::new(Engine::Vortex, Format::OnDiskVortex)])] + #[case::no_targets(&[])] + fn render_table_without_a_baseline_is_a_noop(#[case] targets: &[Target]) -> anyhow::Result<()> { + let mut out = Vec::new(); + render_table( + &mut out, + Vec::::new(), + targets, + )?; + assert!(out.is_empty()); + Ok(()) + } +}