[BUG] dask.uniform_neighbor_sample gives different results depending on the number of GPUs used

[VibhuJawa]

Describe the bug

dask.uniform_neighbor_sample gives different results depending on the number of GPUs used. Results for a single GPU are correct while others seem to be incorrect

See below MRE

Correct

Number of workers = 1, Fanout 1, Number of Samples: 1
Number of workers = 1, Fanout 2, Number of Samples: 2
Number of workers = 1, Fanout 3, Number of Samples: 3

Number of workers = 2, Fanout 1, Number of Samples: 2
Number of workers = 2, Fanout 2, Number of Samples: 4
Number of workers = 2, Fanout 3, Number of Samples: 5

Number of workers = 3, Fanout 1, Number of Samples: 3
Number of workers = 3, Fanout 2, Number of Samples: 5
Number of workers = 3, Fanout 3, Number of Samples: 7

Number of workers = 4, Fanout 1, Number of Samples: 4
Number of workers = 4, Fanout 2, Number of Samples: 7
Number of workers = 4, Fanout 3, Number of Samples: 10

Steps/Code to reproduce bug

MRE:

def run_multigraph_sampling(client):
    import cugraph.dask.comms.comms as Comms
    Comms.initialize(p2p=True)
    
    n_workers = len(client.scheduler_info()['workers'])
    import cudf
    import dask_cudf
    import numpy as np

    import cugraph
    df = cudf.DataFrame({'src':[6]*128,
                         'dst':np.arange(0,128)
                        })
    df = df.astype('int32')

    dask_df = dask_cudf.from_cudf(df, npartitions=128)

    mg_G = cugraph.MultiGraph(directed=True)
    mg_G.from_dask_cudf_edgelist(dask_df,
                                 source='src',
                                 destination='dst',
                                 renumber=True,
                                 legacy_renum_only=True)

    for fanout in [1,2,3]:
        output_df = cugraph.dask.uniform_neighbor_sample(
            mg_G,
            cudf.Series([6]).astype('int32'),
            fanout_vals=[fanout],
            with_replacement=False,
        )
        print(f"Number of workers = {n_workers}, Fanout {fanout}, Number of Samples: {len(output_df)}")
    
    Comms.destroy()

Cluster Stup Code

from dask_cuda import LocalCUDACluster
from dask.distributed import Client


for CUDA_VISIBLE_DEVICES in ['4', '4,5', '4,5,6', '4,5,6,7']:
    with LocalCUDACluster(protocol='tcp',
                               local_directory='/raid/vjawa/dask-cuda-dir/',
                               CUDA_VISIBLE_DEVICES=CUDA_VISIBLE_DEVICES
                           ) as cluster, Client(cluster) as client:
        
        
        
        run_multigraph_sampling(client)
        
        client.shutdown()

Expected behavior

I would expect correct results indepent of the number of GPUS being used.

Environment details

# packages in environment at /datasets/vjawa/miniconda3/envs/cugraph_dev_sep_29:
cugraph                   22.10.0a0+88.g91598080          pypi_0    pypi
libcugraphops             22.10.00a220929 cuda11_g553bacf_29    rapidsai-nightly
pylibcugraph              22.10.0a0+88.g91598080          pypi_0    pypi

Additional context

Probably related Issue: #2760

Probably related PR that can fix this: #2751

CC: @ChuckHastings
CC: @alexbarghi-nv , @rlratzel

[ChuckHastings]

Will recheck when PR 2751 is ready. Thanks for reporting so I can check this condition as well.

[ChuckHastings]

Same analysis as #2760

[rlratzel]

closed by PR #2765