rapidsai / cudf

cuDF - GPU DataFrame Library
https://docs.rapids.ai/api/cudf/stable/
Apache License 2.0
8.42k stars 900 forks source link

[FEA] cudf::column needs a set_stream() function. #12093

Open nvdbaranec opened 1 year ago

nvdbaranec commented 1 year ago

cudf::column uses rmm::device_buffer for internal storage. rmm::device_buffer internally stores the stream it was created on. This can create issues when creating columns on temporary worker-style streams and then returning them to a primary stream. When the column gets destroyed, rmm throws exceptions about invalid device contexts.

A set_stream(rmm::cuda_stream_view) function that recurses through all children would make it easier to hand off columns between streams.

We initially encountered this issue with cudf::io::column_buffer which has a similar problem but it will also be an issue for columns themselves.

bdice commented 1 year ago

The rmm::device_buffer has a set_stream method. This proposed API would call that method of the underlying buffer.

This seems reasonable from what I can tell, but its utility (transferring ownership between streams) will be limited until we expose streams publicly, aside from the internal I/O methods that @hyperbolic2346 is working on. Perhaps we should wait to add this to the public API until we have a broader discussion of adding streams to libcudf APIs?

nvdbaranec commented 1 year ago

I think that's totally fine. The stuff Mike is working on is just experimental stuff as is, and doesn't actually require this from column itself.

hyperbolic2346 commented 1 year ago

I added a PR showing what we are talking about here. I don't know if this is the best way to expose this, but it hopefully clears it up enough to have a solid conversation about it.