code: def create_mg_graph(graph_data):
      """
      Create a graph instance based on the data to be loaded/generated.
      """
      # 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")
  
      cluster = LocalCUDACluster(
          protocol="ucx",
          rmm_pool_size="25GB",
          CUDA_VISIBLE_DEVICES=visible_devices,
      )
      client = Client(cluster)
      Comms.initialize(p2p=True)
      rmm.reinitialize(pool_allocator=True)
      # Assume strings are names of datasets in the datasets package
      if isinstance(graph_data, str):
          ds = getattr(datasets, graph_data)
          edgelist_df = ds.get_edgelist()
          # FIXME: edgelist_df should have column names that match the defaults
          # for G.from_cudf_edgelist()
          edgelist_df = dask_cudf.from_cudf(edgelist_df)
  
      # Assume dictionary contains RMAT params
      elif isinstance(graph_data, dict):
          scale = graph_data["scale"]
          num_edges = (2**scale) * graph_data["edgefactor"]
          seed = _seed
          edgelist_df = rmat(
              scale,
              num_edges,
              0.57,  # from Graph500
              0.19,  # from Graph500
              0.19,  # from Graph500
              seed,
              clip_and_flip=False,
              scramble_vertex_ids=False,  # FIXME: need to understand relevance of this
              create_using=None,  # None == return edgelist
              mg=True,
          )
          edgelist_df["weight"] = cp.float32(1)
      else:
          raise TypeError(
              f"graph_data can only be str or dict, got {type(graph_data)}"
          )
  
      num_nodes = max(
          edgelist_df["src"].max().compute(),
          edgelist_df["dst"].max().compute(),
      )
  
      # running into issues with smaller partitions
      edgelist_df = edgelist_df.repartition(
          npartitions=edgelist_df.npartitions * 2
      )
  
      num_nodes_dict = {"_N": num_nodes}
  
      gs = CuGraphStorage(
          num_nodes_dict=num_nodes_dict, single_gpu=False
      )
      gs.add_edge_data(
          edgelist_df,
          node_col_names=["dst", "src"],
          canonical_etype=["_N", "C", "_N"],
      )
      return (gs, client, cluster)
  
imports: from dask_cuda import LocalCUDACluster
  from dask.distributed import Client
  import dask_cudf
  import time
  import os
  import pytest
  import numpy
  import cupy
  import pytest_benchmark
  from cugraph.generators import rmat
  from cugraph.experimental import datasets
  from cugraph_benchmarking import params
  from cugraph_dgl import CuGraphStorage
  from cugraph.dask.comms import comms
  import dgl
  import torch
  import rmm
label_desc: Use image node_id 2 for calling a global function with example usage: create_mg_graph(graph_data) and returns: gs,  client,  cluster