From 8b02e241617df8bac33d0fd69b03046d2ddaf2d9 Mon Sep 17 00:00:00 2001
From: Naim <110031745+naimnv@users.noreply.github.com>
Date: Tue, 26 Sep 2023 15:43:02 +0200
Subject: [PATCH] Enable temporarily disabled MG tests (#3837)

Enable TEMPORARILY disable single-GPU "MG" tests
And Skip deleting copied Dataframe while creating distributed graph from cudf edge-lists.
Ideally we would like to merger this PR once the [issue 3790](https://github.com/rapidsai/cugraph/issues/3790) is closed, but me  might need to merger it if the issue is not resolved before the next release.

Authors:
  - Naim (https://github.com/naimnv)

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)
  - AJ Schmidt (https://github.com/ajschmidt8)

URL: https://github.com/rapidsai/cugraph/pull/3837
---
 ci/test_python.sh                                           | 6 +-----
 ci/test_wheel.sh                                            | 4 +---
 .../graph_implementation/simpleDistributedGraph.py          | 5 ++++-
 python/cugraph/cugraph/tests/traversal/test_bfs_mg.py       | 5 ++++-
 4 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/ci/test_python.sh b/ci/test_python.sh
index 14886909fc9..7b0077991ae 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -63,10 +63,6 @@ pytest \
   tests
 popd
 
-# FIXME: TEMPORARILY disable single-GPU "MG" testing until
-# https://github.com/rapidsai/cugraph/issues/3790 is closed
-# When closed, replace -k "not _mg" with
-#  -k "not test_property_graph_mg" \
 rapids-logger "pytest cugraph"
 pushd python/cugraph/cugraph
 export DASK_WORKER_DEVICES="0"
@@ -79,7 +75,7 @@ pytest \
   --cov=cugraph \
   --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cugraph-coverage.xml" \
   --cov-report=term \
-  -k "not _mg" \
+  -k "not test_property_graph_mg" \
   tests
 popd
 
diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh
index b62635d08b4..146186ae2e7 100755
--- a/ci/test_wheel.sh
+++ b/ci/test_wheel.sh
@@ -18,7 +18,5 @@ arch=$(uname -m)
 if [[ "${arch}" == "aarch64" && ${RAPIDS_BUILD_TYPE} == "pull-request" ]]; then
     python ./ci/wheel_smoke_test_${package_name}.py
 else
-    # FIXME: TEMPORARILY disable single-GPU "MG" testing until
-    # https://github.com/rapidsai/cugraph/issues/3790 is closed
-    RAPIDS_DATASET_ROOT_DIR=`pwd`/datasets python -m pytest -k "not _mg" ./python/${package_name}/${package_name}/tests
+    RAPIDS_DATASET_ROOT_DIR=`pwd`/datasets python -m pytest ./python/${package_name}/${package_name}/tests
 fi
diff --git a/python/cugraph/cugraph/structure/graph_implementation/simpleDistributedGraph.py b/python/cugraph/cugraph/structure/graph_implementation/simpleDistributedGraph.py
index 0586d0d853c..01885c2d1c3 100644
--- a/python/cugraph/cugraph/structure/graph_implementation/simpleDistributedGraph.py
+++ b/python/cugraph/cugraph/structure/graph_implementation/simpleDistributedGraph.py
@@ -14,6 +14,7 @@
 import gc
 from typing import Union
 import warnings
+import random
 
 import cudf
 import cupy as cp
@@ -182,7 +183,9 @@ def __from_edgelist(
         # Repartition to 2 partitions per GPU for memory efficient process
         input_ddf = input_ddf.repartition(npartitions=len(workers) * 2)
         # FIXME: Make a copy of the input ddf before implicitly altering it.
-        input_ddf = input_ddf.map_partitions(lambda df: df.copy())
+        input_ddf = input_ddf.map_partitions(
+            lambda df: df.copy(), token="custom-" + str(random.random())
+        )
         # The dataframe will be symmetrized iff the graph is undirected
         # otherwise, the inital dataframe will be returned
         if edge_attr is not None:
diff --git a/python/cugraph/cugraph/tests/traversal/test_bfs_mg.py b/python/cugraph/cugraph/tests/traversal/test_bfs_mg.py
index 8ffbecea4fc..5eafc231141 100644
--- a/python/cugraph/cugraph/tests/traversal/test_bfs_mg.py
+++ b/python/cugraph/cugraph/tests/traversal/test_bfs_mg.py
@@ -12,6 +12,7 @@
 # limitations under the License.
 
 import gc
+import random
 
 import pytest
 
@@ -61,7 +62,9 @@ def modify_dataset(df):
         return cudf.concat([df, temp_df])
 
     meta = ddf._meta
-    ddf = ddf.map_partitions(modify_dataset, meta=meta)
+    ddf = ddf.map_partitions(
+        modify_dataset, meta=meta, token="custom-" + str(random.random())
+    )
 
     df = cudf.read_csv(
         input_data_path,