WIP - Add tests for UCX and cuDF serialization

distributed/comm/tests/test_ucx.py

@@ -10,7 +10,7 @@
 from distributed.comm import ucx, parse_address
 from distributed.protocol import to_serialize
 from distributed.deploy.local import LocalCluster
-from distributed.utils_test import gen_test, loop, inc  # noqa: 401
+from distributed.utils_test import gen_cluster, gen_test, loop, inc  # noqa: 401
 
 from .test_comms import check_deserialize
 
@@ -294,3 +294,20 @@ def test_tcp_localcluster(loop):
         #     assert any(w.data == {x.key: 2} for w in c.workers)
         #     assert e.loop is c.loop
         #     print(c.scheduler.workers)
+
+
+@pytest.mark.asyncio
+async def test_cudf_join():
+    from dask.distributed import Scheduler, Worker
+    import dask
+
+    cudf = pytest.importorskip("cudf")
+    async with Scheduler(protocol="ucx", port=0, interface="ib0") as s:
+        async with Worker(s.address, port=0) as a, Worker(s.address, port=0) as b:
+            async with Client(s.address, asynchronous=True) as c:
+                df = dask.datasets.timeseries(
+                    dtypes={"x": int, "y": float}, freq="1s"
+                ).partitions[:2]
+                df = df.map_partitions(cudf.from_pandas)
+                await c.compute(df.x.sum())
+                await c.compute(df[["x"]].sum())

distributed/protocol/cudf.py

@@ -34,7 +34,6 @@ def serialize_cudf_dataframe(x):
     arrays.extend(null_masks)
 
     header = {
-        "is_cuda": len(arrays),
         "subheaders": sub_headers,
         # TODO: the header must be msgpack (de)serializable.
         # See if we can avoid names, and just use integer positions.

distributed/protocol/cupy.py

@@ -24,7 +24,6 @@ def serialize_cupy_ndarray(x):
     # used in the ucx comms for gpu/cpu message passing
     # 'lengths' set by dask
     header = x.__cuda_array_interface__.copy()
-    header["is_cuda"] = 1
     header["dtype"] = dtype
     return header, [data]
 

distributed/protocol/numba.py

@@ -29,7 +29,6 @@ def serialize_numba_ndarray(x):
     # used in the ucx comms for gpu/cpu message passing
     # 'lengths' set by dask
     header = x.__cuda_array_interface__.copy()
-    header["is_cuda"] = 1
     header["dtype"] = dtype
     return header, [data]
 

distributed/protocol/tests/test_cudf.py

@@ -0,0 +1,29 @@
+import pytest
+
+cudf = pytest.importorskip("cudf")
+
+from distributed.protocol import serialize, deserialize
+import dask.dataframe as dd
+
+
+@pytest.mark.parametrize(
+    "df",
+    [
+        cudf.Series([1, 2, 3]),
+        cudf.Series([1, 2, None]),
+        cudf.DataFrame({"x": [1, 2, 3], "y": [1.0, 2.0, 3.0]}),
+        cudf.DataFrame({"x": [1, 2, 3], "s": ["a", "bb", "ccc"]}),
+        cudf.DataFrame(
+            {"x": [1, 2, None], "y": [1.0, 2.0, None], "s": ["a", "bb", None]}
+        ),
+    ],
+)
+def test_basic(df):
+    header, frames = serialize(
+        df, serializers=("cuda", "dask", "pickle"), on_error="raise"
+    )
+    assert header["serializer"] == "cuda"
+    assert not any(isinstance(frame, (bytes, memoryview)) for frame in frames)
+
+    df2 = deserialize(header, frames, deserializers=("cuda", "dask", "pickle"))
+    dd.assert_eq(df, df2)