Skip to content

Commit

Permalink
Re-configure benchmarking devices & add markers to `bench_cugraph_uni…
Browse files Browse the repository at this point in the history
…form_neighbor_sample` (#4561)

This PR re-enables the `benchmarks/cugraph/pytest-based/bench_cugraph_uniform_neighbor_sample.py` benchmark in the MNMG nightlies. 

This benchmark was previously being skipped due to the fact that the benchmarks are configured to run with Pytest markers `"managedmem_off and poolallocator_on"`. Thanks to @jameslamb for spotting this inside `conftest.py`

By adding the missing markers to the test (and removing some `dgx` machine specific dask-client configs), the benchmark should properly run in the nightly jobs.

Authors:
  - Ralph Liu (https://github.com/nv-rliu)

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)
  - James Lamb (https://github.com/jameslamb)

URL: #4561
  • Loading branch information
nv-rliu authored Aug 5, 2024
1 parent b64bc3e commit 1be81a4
Showing 1 changed file with 8 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,11 @@
_seed = 42


###############################################################################
# Helpers
###############################################################################


def create_graph(graph_data):
"""
Create a graph instance based on the data to be loaded/generated, return a
Expand Down Expand Up @@ -107,17 +112,13 @@ def create_mg_graph(graph_data):
Create a graph instance based on the data to be loaded/generated, return a
tuple containing (graph_obj, num_verts, client, cluster)
"""
# range starts at 1 to let let 0 be used by benchmark/client process
visible_devices = os.getenv("DASK_WORKER_DEVICES", "1,2,3,4")

(client, cluster) = start_dask_client(
# enable_tcp_over_ucx=True,
# enable_infiniband=False,
# enable_nvlink=True,
# enable_rdmacm=False,
protocol="ucx",
rmm_pool_size="28GB",
dask_worker_devices=visible_devices,
)
rmm.reinitialize(pool_allocator=True)

Expand Down Expand Up @@ -261,6 +262,9 @@ def uns_func(*args, **kwargs):

################################################################################
# Benchmarks
###############################################################################
@pytest.mark.managedmem_off
@pytest.mark.poolallocator_on
@pytest.mark.parametrize("batch_size", params.batch_sizes.values())
@pytest.mark.parametrize("fanout", [params.fanout_10_25, params.fanout_5_10_15])
@pytest.mark.parametrize(
Expand Down

0 comments on commit 1be81a4

Please sign in to comment.