Skip to content

Run queries on GPUs

Remote queries can run on NVIDIA GPUs with cudf-polars. GPU support is unstable.

Set up

Set instance_type to an NVIDIA GPU instance type, because a type picked from cpus and memory never has a GPU:

ctx = pn.ComputeContext(
    workspace="My workspace",
    instance_type="g4dn.8xlarge",
    cluster_size=2,
)

The instance type must be x86_64 with whole GPUs of compute capability 7.5 or higher, which rules out the p2, g3, p3, g5g and g6f families. GPU clusters also need:

  • Python 3.12 or later on the client, because the cluster uses the client's Python version
  • direct mode, as proxy mode is not supported
  • a single node (cluster_size=1) on instance types with more than one network card, such as p4d.24xlarge and p5.48xlarge
  • CloudFormation template version 1.2.0 or later. When a workspace runs an older one, the portal shows a Workspace update available banner with an Update version link to the stack update in the AWS console.

GPU families have their own EC2 vCPU quota, which can be too low for a first GPU cluster. See the quota error.

In a Python 3.12 or later environment on every worker, install the extra that matches the worker's CUDA major version, cudf-polars-cu12 or cudf-polars-cu13. Install only one: both provide the same cudf_polars module, and installing both can leave a build that does not match the CUDA driver. Until cudf-polars 26.12 is released, the extras resolve to a RAPIDS nightly, which needs the nightly index:

uv pip install --extra-index-url https://pypi.anaconda.org/rapidsai-wheels-nightly/simple \
    "polars-nanook-compute[cudf-polars-cu12]"

Enable GPU support on every worker of the cluster, in the [worker] section of its config file:

[worker.extras]
cudf_polars.enabled = true

A worker with GPU support enabled uses every GPU it sees. To restrict it, set CUDA_VISIBLE_DEVICES to device indexes, such as 0,2. The worker refuses to start when the extra is missing, when it sees no GPU, or when CUDA_VISIBLE_DEVICES names a device by UUID.

Run a GPU query

Pass engine="gpu" to .remote() and sink the result to a directory (a path without a file extension):

lf.remote(ctx, engine="gpu").sink_parquet(uri="<DIRECTORY_URI>").await_result()
  • Collecting the result and Iceberg sinks are not supported.
  • A GPU query never falls back to the CPU: it fails without GPUs or on an operation cudf-polars does not support.
  • lf.remote(ctx, engine="gpu").single_node() runs the query on the GPUs of one worker instead of distributing it across workers.