Dask-cudf multi partition merge slows down with `ucx`

[VibhuJawa]

Dask-cudf multi partition merge slows down with ucx .

Dask-cudf merge seems to slow down with ucx .

Wall time: (15.4 seconds on tcp) vs (37.8 s on ucx) (exp-01)

In the attached example we see a slow down with ucx vs just using tcp .

Wall Times on exp-01

UCX Time

CPU times: user 19.3 s, sys: 1.97 s, total: 21.2 s
Wall time: 38.4 s
2945293


CPU times: user 16.7 s, sys: 1.71 s, total: 18.4 s
Wall time: 37.8 s
2943379

TCP times

CPU times: user 10.8 s, sys: 815 ms, total: 11.6 s
Wall time: 15.7 s
2944022

CPU times: user 10.9 s, sys: 807 ms, total: 11.7 s
Wall time: 15.4 s
2943697 

Repro Code:

Helper Function to create distributed dask-cudf frame

import dask_cudf
import cudf
import os
import time
import dask.dataframe as dd
import dask.array as da

from dask_cuda import LocalCUDACluster
from dask.distributed import Client,wait
from dask.utils import parse_bytes

def create_random_data(n_rows=1_000,n_parts = 10, n_keys_index_1=100_000,n_keys_index_2=100,n_keys_index_3=100, col_prefix = 'a'):
    
    chunks = n_rows//n_parts

    df = dd.concat([
        da.random.random(n_rows, chunks = chunks).to_dask_dataframe(columns= col_prefix + '_non_merge_1'),
        da.random.random(n_rows, chunks = chunks).to_dask_dataframe(columns= col_prefix + '_non_merge_2'),
        da.random.random(n_rows, chunks = chunks).to_dask_dataframe(columns= col_prefix + '_non_merge_3'),
        da.random.randint(0, n_keys_index_1, size=n_rows,chunks = chunks ).to_dask_dataframe(columns= col_prefix + '_0'),
        da.random.randint(0, n_keys_index_2, size=n_rows, chunks = chunks ).to_dask_dataframe(columns= col_prefix +'_1'),
        da.random.randint(0, n_keys_index_3, size=n_rows, chunks = chunks ).to_dask_dataframe(columns= col_prefix +'_2'),
        
    ], axis=1).persist()
    
    gdf = df.map_partitions(cudf.from_pandas)
    gdf =  gdf.persist()
    _ = wait(gdf)
    return gdf

RMM Setup:

def setup_rmm_pool(client):
    client.run(
        cudf.set_allocator,
        pool=True,
        initial_pool_size= parse_bytes("26GB"),
        allocator="default"
    )
    return None

setup_rmm_pool(client)

Merge Code:

The slow down happens on the merge step.

rows_1, parts_1 = 140_176_770, 245
rows_2, parts_2 = 21_004_393, 171

df_1 = create_random_data(n_rows= rows_1, n_parts = parts_1, col_prefix = 'a')
df_2 = create_random_data(n_rows= rows_2,  n_parts = parts_2, col_prefix = 'b')

merged_df = df_1.merge(df_2, left_on = ['a_0','a_1','a_2'], right_on = ['b_0','b_1','b_2'])
%time len(merged_df)

Additional Context:

There has been discussion about this on our internal slack channel, please see for more context.

[quasiben]

Adding the dask reports for ucx/tcp:

• dask report with ucx

• dask report with tcp

The notable issue in the task graph is that for TCP there is more communication but the workers are fully engaged (no white space):

For UCX, there is less communication workers are not actively working on a particular task (white space):

[jakirkham]

Thanks Vibhu for filing and Ben for including the profiling plot.

It looks like the read time on the worker is taking ~1.5x longer in the UCX case than the TCP case. This may be helped by releasing the GIL during these operations. ( #391 )

We also are spending a lot of time in generating the __cuda_array_interface__ in Numba. Not sure why that is the case. Will look at this more closely.

[beckernick]

We are also spending more time working in the UCX report (3-4x the time). I'm going to be raising a similar issue with a reproducible example for groupby-aggregation behavior, today

[jakirkham]

Ah was looking at the gaps initially based on Ben's comment. Should I be looking at something else?

[beckernick]

Ah was looking at the gaps initially based on Ben's comment. Should I be looking at something else?

I just wanted to highlight the actual time spent doing compute is materially different, which may or may not related but also naively feels like a potential contributor to the slowdown.

[jakirkham]

Meaning the overall runtime? If so, then I think Ben has the right idea to look at gaps. If not, some clarity on the compute time you are looking at would be helpful 🙂

[beckernick]

Yes, sorry for being unclear 😄

From the worker profile in the linked performance report, we spend 130 total seconds doing "compute" with UCX on across the 16 GPUs. We only spend 40 total seconds with TCP across the 16 GPUs.

• 62 seconds doing _concat with UCX, vs 13 seconds with TCP
• 19 vs 9 for shuffle_group

[quasiben]

Another interesting thing to note is that deletion of dependencies takes quite a bit of time in UCX compared with TCP. This can be seen in the Worker Profile (administrative) page:

UCX

TCP

[jakirkham]

No worries. Thanks Nick! Just wanted to make sure we are focusing on the pain point 😉

Yeah that's what I was looking at above, Ben. Agree this is what we should focus on.

[quasiben]

It want to note that these issue could be a symptom of poor work stealing and we are still digging into this idea

[jakirkham]

That may be true.

There is also some cost in building Numba arrays, which we are doing. I've done some work to speed this up. Am still trying to figure out the right way to hook this into what we are doing here.

In particular please see this rough benchmark.

[pentschev]

Here are some timings I gathered with this code:

TCP
Create time: 10.193811178207397
Merge time: 15.630461931228638
Create time: 10.504480838775635
Merge time: 15.304280996322632
Create time: 10.684969425201416
Merge time: 16.014200448989868

UCX (Master):
Create time: 10.919841766357422
Merge time: 22.28807020187378
Create time: 11.757182836532593
Merge time: 23.014573097229004
Create time: 11.073307514190674
Merge time: 22.2984356880188

UCX (v1.7.0):
Create time: 22.65068030357361
Merge time: 45.892472982406616
Create time: 21.901089429855347
Merge time: 44.46129822731018
Create time: 22.363646984100342
Merge time: 44.13629865646362

UCX (v1.7.0 + https://github.com/openucx/ucx/pull/4646):
Create time: 10.750166177749634
Merge time: 23.12590527534485
Create time: 11.505124568939209
Merge time: 22.653675079345703
Create time: 11.365516662597656
Merge time: 20.664494276046753

As we can see above, openucx/ucx#4646 helps (and if cherry-picked to v1.7.0, matches performance of master) but TCP still outperforms UCX.

@quasiben and I explored the __delitem__ that was mentioned before, and just for a rough test we ended up commenting it out just to see if that would affect performance, but it turns out that it doesn't change performance at all.

[quasiben]

In looking more thoroughly through the worker profile there seems to be a fair amount of time managing cuda contexts when using UCX:

@kkraus14 do you happen to have any insights here ?

[jakirkham]

Yeah that's what I was trying to get at with this comment.

Am still trying to find the right way to hook Dask into higher level RMM/dask-cuda. PR ( rapidsai/rmm#264 ) is a start at that.