[BUG]: RuntimeError in `pylibcudf.io.parquet.ChunkedParquetReader.read_chunk` when changing device in cudf-polars

[TomAugspurger]

Describe the bug

The GPU engine for polars allows you to specify a device. Changing the device apparently doesn't mix well with some process-global state used in read_parquet.

Steps/Code to reproduce bug

import polars as pl

fname = 'test.parquet'
data = {"a": [1, 2], "b": [3, 4]}
pl.DataFrame(data).write_parquet(fname)


pl.scan_parquet(fname).collect(engine=pl.GPUEngine(device=0))
print("Device 0 success")
pl.scan_parquet(fname).collect(engine=pl.GPUEngine(device=1))
print("Device 1 success")

which raises with

Device 0 success
Traceback (most recent call last):
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/bug.py", line 12, in <module>
    pl.scan_parquet(fname).collect(engine=pl.GPUEngine(device=1))
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/polars/_utils/deprecation.py", line 97, in wrapper
    return function(*args, **kwargs)
  File "/raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/polars/lazyframe/opt_flags.py", line 330, in wrapper
    return function(*args, **kwargs)
  File "/raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/polars/lazyframe/frame.py", line 2407, in collect
    return wrap_df(ldf.collect(engine, callback))
                   ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^
  File "/raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/polars/_utils/scan.py", line 27, in _execute_from_rust
    return function(with_columns, *args)
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/python/cudf_polars/cudf_polars/callback.py", line 263, in _callback
    return evaluate_streaming(ir, config_options)
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/python/cudf_polars/cudf_polars/experimental/parallel.py", line 284, in evaluate_streaming
    return get_scheduler(config_options)(graph, key).to_polars()
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/python/cudf_polars/cudf_polars/experimental/scheduler.py", line 149, in synchronous_scheduler
    cache[k] = _execute_task(graph[k], cache)
               ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/python/cudf_polars/cudf_polars/experimental/scheduler.py", line 51, in _execute_task
    return arg[0](*(_execute_task(a, cache) for a in arg[1:]))
           ~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/python/cudf_polars/cudf_polars/dsl/ir.py", line 788, in do_evaluate
    chunk = reader.read_chunk()
  File "pylibcudf/io/parquet.pyx", line 366, in pylibcudf.io.parquet.ChunkedParquetReader.read_chunk
  File "pylibcudf/io/parquet.pyx", line 385, in pylibcudf.io.parquet.ChunkedParquetReader.read_chunk
RuntimeError: after determining tmp storage requirements for exclusive_scan: cudaErrorInvalidDevice: invalid device ordinal

Expected behavior

No error.

Additional context

Initially reported at pola-rs/polars#25262 (cc @danielcrane)

The error seems to be in pylibcudf. Just using pl.LazyFrame(data).collect() with the two devices works fine.

[danielcrane]

Thanks @TomAugspurger

[TomAugspurger]

Here's a pylibcudf+RMM reproducer, that roughly mimics what we do in cudf-polars:

import polars as pl
import pylibcudf as plc
import pyarrow as pa
import pyarrow.parquet as pq
import rmm

t = pa.Table.from_pydict({"a": [1, 2], "b": [3, 4]})
pq.write_table(t, "test.parquet")

source = plc.io.SourceInfo(["test.parquet"])
options = plc.io.parquet.ParquetReaderOptions.builder(source).build()

rmm._cuda.gpu.setDevice(0)
rmm.mr.set_current_device_resource(rmm.mr.CudaAsyncMemoryResource())
plc.io.parquet.read_parquet(options)

rmm._cuda.gpu.setDevice(1)
rmm.mr.set_current_device_resource(rmm.mr.CudaAsyncMemoryResource())
plc.io.parquet.read_parquet(options)

which raises

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[5], line 1
----> 1 plc.io.parquet.read_parquet(options)

File /raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/pylibcudf/io/parquet.pyx:390, in pylibcudf.io.parquet.read_parquet()

File /raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/pylibcudf/io/parquet.pyx:413, in pylibcudf.io.parquet.read_parquet()

RuntimeError: after determining tmp storage requirements for exclusive_scan: cudaErrorInvalidDevice: invalid device ordinal

I'm not too familiar with pylibcudf's policy on devices. @vyasr do you know whether this looks like a valid use of pylibucdf (and so something needs fixing in read_parquet), or are we doing something wrong in cudf-polars?

[wence-]

This is bad code in libcudf. We have various pinned host memory resources that are used for allocations, but those are static. If someone swaps the device out, then this host mr is invalid. Although host-pinned memory is accessible on any GPU, we use an RMM pool_memory_resource wrapped around the cudaHostAlloc and to get stream-ordering correct that uses events. But the pool_memory_resource is only valid on the device that was active when it was created.

There are various other static lifetime objects that are tied to the current cuda device that would need to be invalidated when switching devices.

I think it is expected that one can't drive libcudf on multiple devices from within the same process, in which case this use case is not supported today in cudf-polars and we should figure out a way to catch that and raise a useful error message.

[wence-]

I've added a check in #20617 so you would now get a useful error message.

[TomAugspurger]

@danielcrane for now, it looks like we won't be able to support switching devices within the same process. Let us know if that's problematic for your use case and we can figure something out.

And just an FYI, we are working to improve multi-GPU execution of polars queries, since it seems like you have a multi-GPU system :)