[REVIEW] Skip adding edge types, edge weights

python/cugraph-dgl/cugraph_dgl/dataloading/utils/extract_graph_helpers.py

@@ -25,8 +25,13 @@ def create_cugraph_graph_from_edges_dict(
 ):
     if edge_dir == "in":
         edges_dict = {k: reverse_edges(df) for k, df in edges_dict.items()}
-    # TODO: Potentially skip this for memory efficiencies
-    edges_dict = {k: add_etype_id(df, etype_id_dict[k]) for k, df in edges_dict.items()}
+    if len(edges_dict) > 1:
+        has_multiple_etypes = True
+        edges_dict = {
+            k: add_etype_id(df, etype_id_dict[k]) for k, df in edges_dict.items()
+        }
+    else:
+        has_multiple_etypes = False
 
     edges_dfs = list(edges_dict.values())
     del edges_dict
@@ -36,24 +41,26 @@ def create_cugraph_graph_from_edges_dict(
         edges_df = cudf.concat(edges_dfs, ignore_index=True)
     del edges_dfs
 
-    edges_df["wgt"] = np.float32(0)
     G = cugraph.MultiGraph(directed=True)
     if isinstance(edges_df, dask_cudf.DataFrame):
-        G.from_dask_cudf_edgelist(
-            edges_df,
-            source="_SRC_",
-            destination="_DST_",
-            edge_attr=["wgt", "_EDGE_ID_", "etp"],
-            renumber=True,
-        )
+        g_creation_f = G.from_dask_cudf_edgelist
+    else:
+        g_creation_f = G.from_cudf_edgelist
+
+    if has_multiple_etypes:
+        edge_etp = "etp"
     else:
-        G.from_cudf_edgelist(
-            edges_df,
-            source="_SRC_",
-            destination="_DST_",
-            edge_attr=["wgt", "_EDGE_ID_", "etp"],
-            renumber=True,
-        )
+        edge_etp = None
+
+    g_creation_f(
+        edges_df,
+        source="_SRC_",
+        destination="_DST_",
+        weight=None,
+        edge_id="_EDGE_ID_",
+        edge_type=edge_etp,
+        renumber=True,
+    )
     return G
 
 

python/cugraph-dgl/examples/muti_trainer_MG_example/workflow.py

@@ -220,6 +220,11 @@ def run_workflow(rank, devices, scheduler_address):
 
 
 if __name__ == "__main__":
+    # Load dummy first
+    # because new environments
+    # require dataset download
+    _ = load_dgl_dataset()
+    del _
     dask_worker_devices = [5, 6]
     scheduler_address = setup_cluster(dask_worker_devices)
 

python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py

@@ -119,13 +119,6 @@ def convert_to_cudf(cp_arrays, weight_t, with_edge_properties, return_offsets=Fa
         df[edge_type_n] = edge_types
         df[hop_id_n] = hop_ids
 
-        print(
-            f"sources: {sources}\n"
-            f"destinations: {destinations}\n"
-            f"batch: {batch_ids}\n"
-            f"offset: {offsets}\n"
-        )
-
         if return_offsets:
             offsets_df = cudf.DataFrame(
                 {