class_docstr: Sampled graph data loader. Wrap a :class:`~cugraph_dgl.CuGraphStorage` and a
  :class:`~cugraph_dgl.dataloading.NeighborSampler` into
  an iterable over mini-batches of samples. cugraph_dgl's ``DataLoader`` extends
  PyTorch's ``DataLoader`` by handling creation and
  transmission of graph samples.
class_docstr_tok: ["Sampled","graph","data","loader",".","Wrap","a",":","class",":","`","~cugraph_dgl.CuGraphStorage","`","and","a",":","class",":","`","~cugraph_dgl.dataloading.NeighborSampler","`","into","an","iterable","over","mini-batches","of","samples",".","cugraph_dgl","'s","``","DataLoader","``","extends","PyTorch","'s","``","DataLoader","``","by","handling","creation","and","transmission","of","graph","samples","."]
docstr: Constructor for CuGraphStorage:
  -------------------------------
  graph : CuGraphStorage
      The graph.
  indices : Tensor or dict[ntype, Tensor]
      The set of indices.  It can either be a tensor of
      integer indices or a dictionary of types and indices.
      The actual meaning of the indices is defined by the :meth:`sample` method of
      :attr:`graph_sampler`.
  graph_sampler : cugraph_dgl.dataloading.NeighborSampler
      The subgraph sampler.
  sampling_output_dir: str
      Output directory to share sampling results in
  batches_per_partition: int
      The number of batches of sampling results to write/read
  seeds_per_call: int
      The number of seeds to sample at once
  device : device context, optional
      The device of the generated MFGs in each iteration, which should be a
      PyTorch device object (e.g., ``torch.device``).
      By default this returns the tenors on device with the current
      cuda context
  use_ddp : boolean, optional
      If True, tells the DataLoader to split the training set for each
      participating process appropriately using
      :class:`torch.utils.data.distributed.DistributedSampler`.
      Overrides the :attr:`sampler` argument of
      :class:`torch.utils.data.DataLoader`.
  ddp_seed : int, optional
      The seed for shuffling the dataset in
      :class:`torch.utils.data.distributed.DistributedSampler`.
      Only effective when :attr:`use_ddp` is True.
  batch_size: int
      Batch size.
  sparse_format: str, default = "coo"
      The sparse format of the emitted sampled graphs. Choose between "csc"
      and "coo". When using "csc", the graphs are of type
      cugraph_dgl.nn.SparseGraph.
  kwargs : dict
      Key-word arguments to be passed to the parent PyTorch
      :py:class:`torch.utils.data.DataLoader` class. Common arguments are:
          - ``batch_size`` (int): The number of indices in each batch.
          - ``drop_last`` (bool): Whether to drop the last incomplete
                                  batch.
          - ``shuffle`` (bool): Whether to randomly shuffle the
                                indices at each epoch
  Examples
  --------
  To train a 3-layer GNN for node classification on a set of nodes
  ``train_nid`` on a homogeneous graph where each node takes messages
  from 15 neighbors on the first layer, 10 neighbors on the second, and
  5 neighbors on the third:
  >>> sampler = cugraph_dgl.dataloading.NeighborSampler([15, 10, 5])
  >>> dataloader = cugraph_dgl.dataloading.DataLoader(
  ...     g, train_nid, sampler,
  ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=0)
  >>> for input_nodes, output_nodes, blocks in dataloader:
  ...     train_on(input_nodes, output_nodes, blocks)
  **Using with Distributed Data Parallel**
  If you are using PyTorch's distributed training (e.g. when using
  :mod:`torch.nn.parallel.DistributedDataParallel`),
  you can train the model by turning
  on the `use_ddp` option:
  >>> sampler = cugraph_dgl.dataloading.NeighborSampler([15, 10, 5])
  >>> dataloader = cugraph_dgl.dataloading.DataLoader(
  ...     g, train_nid, sampler, use_ddp=True,
  ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=0)
  >>> for epoch in range(start_epoch, n_epochs):
  ...     for input_nodes, output_nodes, blocks in dataloader:
  ...
docstr_tok: ["Constructor","for","CuGraphStorage",":","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","-","graph",":","CuGraphStorage","The","graph",".","indices",":","Tensor","or","dict","[","ntype",",","Tensor","]","The","set","of","indices",".","It","can","either","be","a","tensor","of","integer","indices","or","a","dictionary","of","types","and","indices",".","The","actual","meaning","of","the","indices","is","defined","by","the",":","meth",":","`","sample","`","method","of",":","attr",":","`","graph_sampler","`",".","graph_sampler",":","cugraph_dgl.dataloading.NeighborSampler","The","subgraph","sampler",".","sampling_output_dir",":","str","Output","directory","to","share","sampling","results","in","batches_per_partition",":","int","The","number","of","batches","of","sampling","results","to","write\/read","seeds_per_call",":","int","The","number","of","seeds","to","sample","at","once","device",":","device","context",",","optional","The","device","of","the","generated","MFGs","in","each","iteration",",","which","should","be","a","PyTorch","device","object","(","e.g.",",","``","torch.device","``",")",".","By","default","this","returns","the","tenors","on","device","with","the","current","cuda","context","use_ddp",":","boolean",",","optional","If","True",",","tells","the","DataLoader","to","split","the","training","set","for","each","participating","process","appropriately","using",":","class",":","`","torch.utils.data.distributed.DistributedSampler","`",".","Overrides","the",":","attr",":","`","sampler","`","argument","of",":","class",":","`","torch.utils.data.DataLoader","`",".","ddp_seed",":","int",",","optional","The","seed","for","shuffling","the","dataset","in",":","class",":","`","torch.utils.data.distributed.DistributedSampler","`",".","Only","effective","when",":","attr",":","`","use_ddp","`","is","True",".","batch_size",":","int","Batch","size",".","sparse_format",":","str",",","default","=","``","coo","''","The","sparse","format","of","the","emitted","sampled","graphs",".","Choose","between","``","csc","''","and","``","coo","''",".","When","using","``","csc","''",",","the","graphs","are","of","type","cugraph_dgl.nn.SparseGraph",".","kwargs",":","dict","Key-word","arguments","to","be","passed","to","the","parent","PyTorch",":","py",":","class",":","`","torch.utils.data.DataLoader","`","class",".","Common","arguments","are",":","-","``","batch_size","``","(","int",")",":","The","number","of","indices","in","each","batch",".","-","``","drop_last","``","(","bool",")",":","Whether","to","drop","the","last","incomplete","batch",".","-","``","shuffle","``","(","bool",")",":","Whether","to","randomly","shuffle","the","indices","at","each","epoch","Examples","--","--","--","--","To","train","a","3-layer","GNN","for","node","classification","on","a","set","of","nodes","``","train_nid","``","on","a","homogeneous","graph","where","each","node","takes","messages","from","15","neighbors","on","the","first","layer",",","10","neighbors","on","the","second",",","and","5","neighbors","on","the","third",":",">",">",">","sampler","=","cugraph_dgl.dataloading.NeighborSampler","(","[","15",",","10",",","5","]",")",">",">",">","dataloader","=","cugraph_dgl.dataloading.DataLoader","(","...","g",",","train_nid",",","sampler",",","...","batch_size=1024",",","shuffle=True",",","drop_last=False",",","num_workers=0",")",">",">",">","for","input_nodes",",","output_nodes",",","blocks","in","dataloader",":","...","train_on","(","input_nodes",",","output_nodes",",","blocks",")","*","*","Using","with","Distributed","Data","Parallel","*","*","If","you","are","using","PyTorch","'s","distributed","training","(","e.g",".","when","using",":","mod",":","`","torch.nn.parallel.DistributedDataParallel","`",")",",","you","can","train","the","model","by","turning","on","the","`","use_ddp","`","option",":",">",">",">","sampler","=","cugraph_dgl.dataloading.NeighborSampler","(","[","15",",","10",",","5","]",")",">",">",">","dataloader","=","cugraph_dgl.dataloading.DataLoader","(","...","g",",","train_nid",",","sampler",",","use_ddp=True",",","...","batch_size=1024",",","shuffle=True",",","drop_last=False",",","num_workers=0",")",">",">",">","for","epoch","in","range","(","start_epoch",",","n_epochs",")",":","...","for","input_nodes",",","output_nodes",",","blocks","in","dataloader",":","..."]
code: def __init__(
      self,
      graph: cugraph_dgl.CuGraphStorage,
      indices: torch.Tensor,
      graph_sampler: cugraph_dgl.dataloading.NeighborSampler,
      sampling_output_dir: str,
      batches_per_partition: int = 50,
      seeds_per_call: int = 200_000,
      device: torch.device = None,
      use_ddp: bool = False,
      ddp_seed: int = 0,
      batch_size: int = 1024,
      drop_last: bool = False,
      shuffle: bool = False,
      sparse_format: str = "coo",
      **kwargs,
  ):
      """
      Constructor for CuGraphStorage:
      -------------------------------
      graph : CuGraphStorage
          The graph.
      indices : Tensor or dict[ntype, Tensor]
          The set of indices.  It can either be a tensor of
          integer indices or a dictionary of types and indices.
          The actual meaning of the indices is defined by the :meth:`sample` method of
          :attr:`graph_sampler`.
      graph_sampler : cugraph_dgl.dataloading.NeighborSampler
          The subgraph sampler.
      sampling_output_dir: str
          Output directory to share sampling results in
      batches_per_partition: int
          The number of batches of sampling results to write/read
      seeds_per_call: int
          The number of seeds to sample at once
      device : device context, optional
          The device of the generated MFGs in each iteration, which should be a
          PyTorch device object (e.g., ``torch.device``).
          By default this returns the tenors on device with the current
          cuda context
      use_ddp : boolean, optional
          If True, tells the DataLoader to split the training set for each
          participating process appropriately using
          :class:`torch.utils.data.distributed.DistributedSampler`.
          Overrides the :attr:`sampler` argument of
          :class:`torch.utils.data.DataLoader`.
      ddp_seed : int, optional
          The seed for shuffling the dataset in
          :class:`torch.utils.data.distributed.DistributedSampler`.
          Only effective when :attr:`use_ddp` is True.
      batch_size: int
          Batch size.
      sparse_format: str, default = "coo"
          The sparse format of the emitted sampled graphs. Choose between "csc"
          and "coo". When using "csc", the graphs are of type
          cugraph_dgl.nn.SparseGraph.
      kwargs : dict
          Key-word arguments to be passed to the parent PyTorch
          :py:class:`torch.utils.data.DataLoader` class. Common arguments are:
              - ``batch_size`` (int): The number of indices in each batch.
              - ``drop_last`` (bool): Whether to drop the last incomplete
                                      batch.
              - ``shuffle`` (bool): Whether to randomly shuffle the
                                    indices at each epoch
      Examples
      --------
      To train a 3-layer GNN for node classification on a set of nodes
      ``train_nid`` on a homogeneous graph where each node takes messages
      from 15 neighbors on the first layer, 10 neighbors on the second, and
      5 neighbors on the third:
      >>> sampler = cugraph_dgl.dataloading.NeighborSampler([15, 10, 5])
      >>> dataloader = cugraph_dgl.dataloading.DataLoader(
      ...     g, train_nid, sampler,
      ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=0)
      >>> for input_nodes, output_nodes, blocks in dataloader:
      ...     train_on(input_nodes, output_nodes, blocks)
      **Using with Distributed Data Parallel**
      If you are using PyTorch's distributed training (e.g. when using
      :mod:`torch.nn.parallel.DistributedDataParallel`),
      you can train the model by turning
      on the `use_ddp` option:
      >>> sampler = cugraph_dgl.dataloading.NeighborSampler([15, 10, 5])
      >>> dataloader = cugraph_dgl.dataloading.DataLoader(
      ...     g, train_nid, sampler, use_ddp=True,
      ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=0)
      >>> for epoch in range(start_epoch, n_epochs):
      ...     for input_nodes, output_nodes, blocks in dataloader:
      ...
      """
      if sparse_format not in ["coo", "csc"]:
          raise ValueError(
              f"sparse_format must be one of 'coo', 'csc', "
              f"but got {sparse_format}."
          )
      self.sparse_format = sparse_format
  
      self.ddp_seed = ddp_seed
      self.use_ddp = use_ddp
      self.shuffle = shuffle
      self.drop_last = drop_last
      self.graph_sampler = graph_sampler
      worker_init_fn = dgl.dataloading.WorkerInitWrapper(
          kwargs.get("worker_init_fn", None)
      )
      self.other_storages = {}
      self.epoch_number = 0
      self._batch_size = batch_size
      self._sampling_output_dir = sampling_output_dir
      self._batches_per_partition = batches_per_partition
      self._seeds_per_call = seeds_per_call
      self._rank = None
  
      indices = _dgl_idx_to_cugraph_idx(indices, graph)
  
      self.tensorized_indices_ds = (
          dgl.dataloading.create_tensorized_dataset(
              indices,
              batch_size,
              drop_last,
              use_ddp,
              ddp_seed,
              shuffle,
              kwargs.get("persistent_workers", False),
          )
      )
  
      if len(graph.ntypes) <= 1:
          self.cugraph_dgl_dataset = HomogenousBulkSamplerDataset(
              total_number_of_nodes=graph.total_number_of_nodes,
              edge_dir=self.graph_sampler.edge_dir,
              sparse_format=sparse_format,
          )
      else:
          etype_id_to_etype_str_dict = {
              v: k for k, v in graph._etype_id_dict.items()
          }
  
          self.cugraph_dgl_dataset = HeterogenousBulkSamplerDataset(
              num_nodes_dict=graph.num_nodes_dict,
              etype_id_dict=etype_id_to_etype_str_dict,
              etype_offset_dict=graph._etype_offset_d,
              ntype_offset_dict=graph._ntype_offset_d,
              edge_dir=self.graph_sampler.edge_dir,
          )
  
      if use_ddp:
          rank = torch.distributed.get_rank()
          client = default_client()
          self._graph_creation_event = Event(
              "cugraph_dgl_load_mg_graph_event"
          )
          if rank == 0:
              G = create_cugraph_graph_from_edges_dict(
                  edges_dict=graph._edges_dict,
                  etype_id_dict=graph._etype_id_dict,
                  edge_dir=graph_sampler.edge_dir,
              )
              client.publish_dataset(cugraph_dgl_mg_graph_ds=G)
              self._graph_creation_event.set()
          else:
              if self._graph_creation_event.wait(timeout=1000):
                  G = client.get_dataset("cugraph_dgl_mg_graph_ds")
              else:
                  raise RuntimeError(
                      f"Fetch cugraph_dgl_mg_graph_ds to worker_id {rank}",
                      "from worker_id 0 failed",
                  )
      else:
          rank = 0
          G = create_cugraph_graph_from_edges_dict(
              edges_dict=graph._edges_dict,
              etype_id_dict=graph._etype_id_dict,
              edge_dir=graph_sampler.edge_dir,
          )
  
      self._rank = rank
      self._cugraph_graph = G
      super().__init__(
          self.cugraph_dgl_dataset,
          batch_size=None,
          worker_init_fn=worker_init_fn,
          collate_fn=lambda x: x,  # Hack to prevent collating
          **kwargs,
      )
  
file_path: cugraph/python/cugraph-dgl/cugraph_dgl/dataloading/dataloader.py
imports: from __future__ import annotations
  import os
  import shutil
  import cugraph_dgl
  import cupy
  import cudf
  from cugraph.utilities.utils import import_optional
  from cugraph.gnn import BulkSampler
  from dask.distributed import default_client, Event
  from cugraph_dgl.dataloading import HomogenousBulkSamplerDataset, HeterogenousBulkSamplerDataset
  from cugraph_dgl.dataloading.utils.extract_graph_helpers import create_cugraph_graph_from_edges_dict
label_desc: Use image node_id 1 to create a new DataLoader object from inherited base classes: torch.utils.data with example: obj = DataLoader(graph, indices, graph_sampler, sampling_output_dir, batches_per_partition, seeds_per_call, device, use_ddp, ddp_seed, batch_size, drop_last, shuffle, sparse_format)