From a7b9835fffac3c7ff87202c0e354057eb2d73696 Mon Sep 17 00:00:00 2001 From: Joseph Nke Date: Fri, 3 Dec 2021 12:34:34 -0800 Subject: [PATCH] Shutdown the connected scheduler and workers --- benchmarks/python_e2e/cugraph_dask_funcs.py | 5 ++++- python/cugraph/cugraph/tests/conftest.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/benchmarks/python_e2e/cugraph_dask_funcs.py b/benchmarks/python_e2e/cugraph_dask_funcs.py index c8f48936647..9ff22a9f248 100644 --- a/benchmarks/python_e2e/cugraph_dask_funcs.py +++ b/benchmarks/python_e2e/cugraph_dask_funcs.py @@ -182,6 +182,9 @@ def setup(dask_scheduler_file=None, rmm_pool_size=None): def teardown(client, cluster=None): Comms.destroy() - client.close() + # Shutdown the connected scheduler and workers + # therefore we will no longer rely on killing the dask cluster ID + # for MNMG runs + client.shutdown() if cluster: cluster.close() diff --git a/python/cugraph/cugraph/tests/conftest.py b/python/cugraph/cugraph/tests/conftest.py index 5c507ed28c6..046aa629286 100644 --- a/python/cugraph/cugraph/tests/conftest.py +++ b/python/cugraph/cugraph/tests/conftest.py @@ -58,7 +58,10 @@ def dask_client(): yield client Comms.destroy() - client.close() + # Shut down the connected scheduler and workers + # therefore we will no longer rely on killing the dask cluster ID + # for MNMG runs + client.shutdown() if cluster: cluster.close() print("\ndask_client fixture: client.close() called")