[REVIEW] Optimize dask.uniform_neighbor_sample

[VibhuJawa]

This PR closes #2872 and part of https://github.com/rapidsai/graph_dl/issues/74 .

On a preliminary benchmark i am seeing a 6x , 3.5x improvement and i expect more improvement on bigger clusters.

Before PR :

%%timeit
df = cugraph.dask.uniform_neighbor_sample(dask_g, start_list[:10_000], [20])
187 ms ± 40.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

After PR

%%timeit
df = cugraph.dask.uniform_neighbor_sample(dask_g, start_list[:10_000], [20])
50.1 ms ± 1.39 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Updated benchmark

Before PR:

---------------------------------------------------------------------- benchmark: 3 tests ----------------------------------------------------------------------
Name (time in ms, mem in bytes)                            Mean              GPU mem            GPU Leaked mem            Rounds            GPU Rounds          
----------------------------------------------------------------------------------------------------------------------------------------------------------------
bench_uniform_neigbour_sample_email_eu_core[1000]      170.2833 (1.0)        300,224 (1.0)                   0 (1.0)           1           1
bench_uniform_neigbour_sample_email_eu_core[5000]      191.7691 (1.13)     1,200,944 (4.00)                  0 (1.0)           1           1
bench_uniform_neigbour_sample_email_eu_core[10000]     194.8658 (1.14)     1,200,944 (4.00)                  0 (1.0)           1           1
----------------------------------------------------------------------------------------------------------------------------------------------------------------

After PR:

--------------------------------------------------------------------- benchmark: 3 tests --------------------------------------------------------------------
Name (time in ms, mem in bytes)                           Mean            GPU mem            GPU Leaked mem            Rounds            GPU Rounds          
-------------------------------------------------------------------------------------------------------------------------------------------------------------
bench_uniform_neigbour_sample_email_eu_core[1000]      54.0802 (1.0)            0 (1.0)                   0 (1.0)           1           1
bench_uniform_neigbour_sample_email_eu_core[5000]      63.0271 (1.17)           0 (1.0)                   0 (1.0)           1           1
bench_uniform_neigbour_sample_email_eu_core[10000]     70.1445 (1.30)           0 (1.0)                   0 (1.0)           1           1
-------------------------------------------------------------------------------------------------------------------------------------------------------------

We should probably follow similar things for other algos too.
CC: @jnke2016 , @rlratzel , @alexbarghi-nv .

[codecov-commenter]

Codecov Report

Base: 60.43% // Head: 61.71% // Increases project coverage by +1.27% 🎉

Coverage data is based on head (632d584) compared to base (d86c933).
Patch has no changes to coverable lines.

Additional details and impacted files

@@               Coverage Diff                @@
##           branch-22.12    #2887      +/-   ##
================================================
+ Coverage         60.43%   61.71%   +1.27%     
================================================
  Files               114      123       +9     
  Lines              6547     7543     +996     
================================================
+ Hits               3957     4655     +698     
- Misses             2590     2888     +298     

Impacted Files	Coverage Δ
python/cugraph/cugraph/components/connectivity.py	94.11% <ø> (-0.89%)	⬇️
python/cugraph/cugraph/dask/__init__.py	100.00% <ø> (ø)
...on/cugraph/cugraph/dask/components/connectivity.py	31.03% <ø> (+5.10%)	⬆️
...thon/cugraph/cugraph/dask/sampling/random_walks.py	22.91% <ø> (ø)
...h/cugraph/dask/sampling/uniform_neighbor_sample.py	25.42% <ø> (+4.26%)	⬆️
...ugraph/cugraph/dask/structure/mg_property_graph.py	14.12% <ø> (+0.46%)	⬆️
python/cugraph/cugraph/dask/traversal/bfs.py	25.39% <ø> (-0.84%)	⬇️
...ph/cugraph/experimental/pyg_extensions/__init__.py	100.00% <ø> (ø)
python/cugraph/cugraph/gnn/__init__.py	100.00% <ø> (ø)
...h/cugraph/gnn/dgl_extensions/base_cugraph_store.py	84.21% <ø> (ø)
... and 47 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[alexbarghi-nv]

Just one comment; it would be nice to have a helper function for this so we can apply it to more algorithms. But I suggest we leave that to a future PR given how much work we have left before burndown.

To clarify, this was referencing lines 174-180

[alexbarghi-nv]

👍

[jnke2016]

This implementation experiences a hang on 2, 3 and 4 GPUs similar to this issue that was closed few months ago by this PR. To reproduce, just run this implementation of uniform_nwighbor_sample several times(less than 10) on 2 GPUs.

[VibhuJawa]

This implementation experiences a hang on 2, 3 and 4 GPUs similar to this issue that was closed few months ago by this PR. To reproduce, just run this implementation of uniform_nwighbor_sample several times(less than 10) on 2 GPUs.

Thanks for testing this thoroughly @jnke2016 , I really appreciate it. Can you post a link to the code that you ran to get to a hand as i could not get it to hand in my testing on 8 GPUs .

We can probably add those statements back and give back 10 ms of perf on 8GPUs but it is not a real solution as it will be much worse on larger cluster (say 128 GPUs) .

I am gonna try looking into a real fix of the problem today

[jnke2016]

You can run this on 2 GPUs.

import cugraph
from dask.distributed import Client
from dask_cuda import LocalCUDACluster
import dask
from cugraph.dask.comms import comms as Comms
import cudf
import dask_cudf
from cugraph.dask.common.mg_utils import get_visible_devices
import cugraph.dask as dcg
from cugraph.generators import rmat

import time

def setup():
    cluster = LocalCUDACluster()
    client = Client(cluster)
    client.wait_for_workers(len(get_visible_devices()))

    Comms.initialize(p2p=True)
    return (client, cluster)

def teardown(client, cluster=None):
    Comms.destroy()
    client.close()
    if cluster:
        cluster.close()

def generate_edgelist(scale,
                      edgefactor,
                      seed=None,
                      unweighted=False,
                     ):

    df = rmat(
        scale,
        (2**scale)*edgefactor,
        0.57,  # from Graph500
        0.19,  # from Graph500
        0.19,  # from Graph500
        seed or 42,
        clip_and_flip=False,
        scramble_vertex_ids=True,
        create_using=None,  # return edgelist instead of Graph instance
        mg=True
    )

    return df

if __name__ == "__main__":

    setup_objs = setup()
    client = setup_objs[0]
    num_workers = len(client.scheduler_info()['workers'])

    df = generate_edgelist(scale=10, edgefactor=16, seed=5)

    m_G = cugraph.Graph()
    m_G.from_dask_cudf_edgelist(
        df, source='src', destination='dst', renumber=True, legacy_renum_only=True)

    vertices = m_G.nodes().persist()
    print("number of nodes are ", len(vertices))
    print("number of edges is ", len(m_G.input_df))
    vertices = vertices.compute().reset_index(drop=True).iloc[:10000]


    time_list = []
    for i in range(10):
        print("*****iteration ", i, "*****")
        t0 = time.time()
        results = cugraph.dask.uniform_neighbor_sample(m_G, vertices, [2])
        t1 = time.time()
        total = t1-t0
        time_list.append(total)


    import statistics
    print("the mean is ", statistics.mean(time_list))
    print("the standard deviation is ", statistics.stdev(time_list))
    print("number of GPUs is ", num_workers)
    
    teardown(*setup_objs)

[VibhuJawa]

@jnke2016 , Thanks for this script. Let me investigate more :-)

[VibhuJawa]

@jnke2016 ,

Swapping client.submit to client.map is runnning into issues with work-stealing (See internal thread). Swapping it back gets rid of the hang .

We still see speedups (See updated benchmarks in the PR description on a 8 GPU cluster) but we will not be as fast as we should be on large clusters . (We currently loose 5ms per worker)

Anyways, I have verified that your script works on 1,2,3,4,5,6,7,8 GPUs with 50 iterations each without hanging.

I will make the cluster agnostic changes in 23.02 as it is more involved.

Please feel free to review again

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[VibhuJawa]

@rerun tests

[wangxiaoyunNV]

It looks good.

[jnke2016]

This looks good to me.

[BradReesWork]

@gpucibot merge