WIP - Add tests for UCX and cuDF serialization

[mrocklin]

Currently we don't robustly serialize cudf dataframes or series. This causes issues when trying to communciate them in dask workloads. So far this PR adds a few failing tests around cudf serialization and a larger full dask + cudf workload.

The first thing to do here is probably to extend the serialization functions in distributed/protocol/cudf.py to include cudf.Series as well as less standard cudf issues like text data and missing values.

@kkraus14 any suggestions you may have on serializing cudf dataframes would be very welcome.

@rjzamora this task may interest you

[kkraus14]

@kkraus14 any suggestions you may have on serializing cudf dataframes would be very welcome.

For cudf.Series at its core is 2 Numba devicearrays, one for the data, one for the index, and then an optional 3rd Numba devicearray if there's nulls. The rest of the the info of things like the dtype, null_count, etc. can be shipped in the header pretty easily. It looks like you've handled it minus the index in https://github.com/dask/distributed/blob/master/distributed/protocol/cudf.py#L22 already.

For cudf.DataFrame it looks like there's a bug where the definition should be deserialize_cudf_dataframe (https://github.com/dask/distributed/blob/master/distributed/protocol/cudf.py#L50) but otherwise looks pretty good to me.

[kkraus14]

As far as text data goes, NVStrings has functions to convert the data into Arrow format which would break down into a data buffer, an offsets buffer, and a null buffer (if needed), which you should be able to ship around via your Numba devicearray code. I'm not exactly clear on how we'd plumb that through to the distributed protocol here since it would return those by copy instead of reference, but I imagine it's doable.

[rjzamora]

I'm currently looking into this... As far as I can tell, cudf.Series is a relatively simple extension (it is easy to get those cases to pass for the new test).

I'm a bit more confused by the strings - @kkraus14, can you say a bit more about the correct way to convert the nvstrings data into numba-serializable buffers? Should I be using to_offsets(values, offsets) ?

[mrocklin]

(it is easy to get those cases to pass for the new test).

@rjzamora if you have code for this could I ask you to publish it as a small PR? This would unblock me on some join work.

[rjzamora]

Sorry - I got delayed on this and also broke some of my progress from yesterday. I submitted a rough PR in case it helps.

[mrocklin]

I recently discovered that cudf already has some code for this here:

https://github.com/rapidsai/cudf/blob/7312d4aa413eafdad28b67efccbc3dc15d9b5e27/python/cudf/dataframe/dataframe.py#L122-L135

https://github.com/rapidsai/cudf/blob/7312d4aa413eafdad28b67efccbc3dc15d9b5e27/python/cudf/dataframe/series.py#L102-L111

I wonder if we can do this work there.

[rjzamora]

The cudf component of this PR should be addressed by rapidsai/cudf#1947. However, maybe we still want to incorperate UCX/cuDF serialization testing?

[mrocklin]

That sounds reasonable to me