From d041dddb4c6d06122d47a27403ec6af25fa5d22d Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 4 Mar 2025 07:19:17 -0800 Subject: [PATCH] Document `frames_to_gpumemoryview` --- .../experimental/dask_serialize.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/python/cudf_polars/cudf_polars/experimental/dask_serialize.py b/python/cudf_polars/cudf_polars/experimental/dask_serialize.py index ca34719479e..2144d9480c3 100644 --- a/python/cudf_polars/cudf_polars/experimental/dask_serialize.py +++ b/python/cudf_polars/cudf_polars/experimental/dask_serialize.py @@ -18,6 +18,25 @@ def frames_to_gpumemoryview(frames): + """ + Convert the elements of `frames` to gpumemoryview. + + UCX transfers produce `rmm.DeviceBuffer` objects instead of `gpumemoryview`. + This function leverages CUDA array interface to convert the elements of + `frames` to `gpumemoryview`, if necessary. + + Parameters + ---------- + frames: list[Any] + List of frames to convert to `gpumemoryview` if they implement CUDA + array interface. + + Returns + ------- + converted: list[Any] + List of frames where all frames implementing CUDA array interface have + been converted to `plc.gpumemoryview`. + """ return [ plc.gpumemoryview(f) if hasattr(f, "__cuda_array_interface__") else f for f in frames