graphlearn_torch/python/distributed/dist_dataset.py

# Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, eithPer express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from multiprocessing.reduction import ForkingPickler
from typing import Dict, List, Optional, Union, Literal, Tuple, Callable

import torch

from ..data import Dataset, Graph, Feature, DeviceGroup, vineyard_utils
from ..partition import (
  load_partition, cat_feature_cache,
  PartitionBook, HeteroNodePartitionDict, HeteroEdgePartitionDict
)
from ..typing import (NodeType, EdgeType, NodeLabel, NodeIndex)
from ..utils import share_memory, default_id_filter, default_id_select


class DistDataset(Dataset):
  r""" Graph and feature dataset with distributed partition info.

  """
  def __init__(
    self,
    num_partitions: int = 1,
    partition_idx: int = 0,
    graph_partition: Union[Graph, Dict[EdgeType, Graph]] = None,
    node_feature_partition: Union[Feature, Dict[NodeType, Feature]] = None,
    edge_feature_partition: Union[Feature, Dict[EdgeType, Feature]] = None,
    whole_node_labels: NodeLabel = None,
    node_pb: Union[PartitionBook, HeteroNodePartitionDict] = None,
    edge_pb: Union[PartitionBook, HeteroEdgePartitionDict] = None,
    node_feat_pb: Union[PartitionBook, HeteroNodePartitionDict] = None,
    edge_feat_pb: Union[PartitionBook, HeteroEdgePartitionDict] = None,
    edge_dir: Literal['in', 'out'] = 'out',
    graph_caching: bool = False,
    node_split: Tuple[NodeIndex, NodeIndex, NodeIndex] = None,
    id_filter: Callable = default_id_filter,
    id_select: Callable = default_id_select
  ):
    super().__init__(
      graph_partition,
      node_feature_partition,
      edge_feature_partition,
      whole_node_labels,
      edge_dir,
      node_split,
    )
    self.id_filter = id_filter
    self.id_select = id_select
    self.num_partitions = num_partitions
    self.partition_idx = partition_idx
    self.graph_caching = graph_caching

    self.node_pb = node_pb
    self.edge_pb = edge_pb

    # As the loaded feature partition may be concatenated with its cached
    # features and the partition book for features will be modified, thus we
    # need to distinguish them with the original graph partition books.
    #
    # If the `node_feat_pb` or `edge_feat_pb` is not provided, the `node_pb`
    # or `edge_pb` will be used instead for feature lookups.
    self._node_feat_pb = node_feat_pb
    self._edge_feat_pb = edge_feat_pb

    if self.graph is not None:
      assert self.node_pb is not None
    if self.node_features is not None:
      assert self.node_pb is not None or self._node_feat_pb is not None
    if self.edge_features is not None:
      assert self.edge_pb is not None or self._edge_feat_pb is not None

  def load(
    self,
    root_dir: str,
    partition_idx: int,
    graph_mode: str = 'ZERO_COPY',
    input_layout: Literal['COO', 'CSR', 'CSC'] = 'COO',
    feature_with_gpu: bool = True,
    graph_caching: bool = False,
    device_group_list: Optional[List[DeviceGroup]] = None,
    whole_node_label_file: Union[str, Dict[NodeType, str]] = None,
    device: Optional[int] = None
  ):
    r""" Load a certain dataset partition from partitioned files and create
    in-memory objects (``Graph``, ``Feature`` or ``torch.Tensor``).

    Args:
      root_dir (str): The directory path to load the graph and feature
        partition data.
      partition_idx (int): Partition idx to load.
      graph_mode (str): Mode for creating graphlearn_torch's ``Graph``, including
        ``CPU``, ``ZERO_COPY`` or ``CUDA``. (default: ``ZERO_COPY``)
      input_layout (str): layout of the input graph, including ``CSR``, ``CSC`` 
        or ``COO``. (default: ``COO``)
      feature_with_gpu (bool): A Boolean value indicating whether the created
        ``Feature`` objects of node/edge features use ``UnifiedTensor``.
        If True, it means ``Feature`` consists of ``UnifiedTensor``, otherwise
        ``Feature`` is a PyTorch CPU Tensor, the ``device_group_list`` and
        ``device`` will be invliad. (default: ``True``)
      graph_caching (bool): A Boolean value indicating whether to load the full
        graph totoploy instead of partitioned one.
      device_group_list (List[DeviceGroup], optional): A list of device groups
        used for feature lookups, the GPU part of feature data will be
        replicated on each device group in this list during the initialization.
        GPUs with peer-to-peer access to each other should be set in the same
        device group properly.  (default: ``None``)
      whole_node_label_file (str): The path to the whole node labels which are
        not partitioned. (default: ``None``)
      device: The target cuda device rank used for graph operations when graph
        mode is not "CPU" and feature lookups when the GPU part is not None.
        (default: ``None``)
    """
    (
      self.num_partitions,
      self.partition_idx,
      graph_data,
      node_feat_data,
      edge_feat_data,
      self.node_pb,
      self.edge_pb
    ) = load_partition(root_dir, partition_idx, graph_caching)

    # init graph partition
    if isinstance(graph_data, dict):
      # heterogeneous.
      edge_index, edge_ids, edge_weights = {}, {}, {}
      for k, v in graph_data.items():
        edge_index[k] = v.edge_index
        edge_ids[k] = v.eids
        edge_weights[k] = v.weights
    else:
      # homogeneous.
      edge_index = graph_data.edge_index
      edge_ids = graph_data.eids
      edge_weights = graph_data.weights
    self.init_graph(edge_index, edge_ids, edge_weights, layout=input_layout,
                    graph_mode=graph_mode, device=device)
    self.graph_caching = graph_caching

    # load node feature partition
    if node_feat_data is not None:
      node_cache_ratio, node_feat, node_feat_id2idx, node_feat_pb = \
        _cat_feature_cache(partition_idx, node_feat_data, self.node_pb)
      self.init_node_features(
        node_feat, node_feat_id2idx, None, node_cache_ratio,
        device_group_list, device, feature_with_gpu, dtype=None
      )
      self._node_feat_pb = node_feat_pb

    # load edge feature partition
    if edge_feat_data is not None:
      edge_cache_ratio, edge_feat, edge_feat_id2idx, edge_feat_pb = \
        _cat_feature_cache(partition_idx, edge_feat_data, self.edge_pb)
      self.init_edge_features(
        edge_feat, edge_feat_id2idx, edge_cache_ratio,
        device_group_list, device, feature_with_gpu, dtype=None
      )
      self._edge_feat_pb = edge_feat_pb

    # load whole node labels
    if whole_node_label_file is not None:
      if isinstance(whole_node_label_file, dict):
        whole_node_labels = {}
        for ntype, file in whole_node_label_file.items():
          whole_node_labels[ntype] = torch.load(file)
      else:
        whole_node_labels = torch.load(whole_node_label_file)
      self.init_node_labels(whole_node_labels)

  def random_node_split(
    self,
    num_val: Union[float, int],
    num_test: Union[float, int],
  ):
    r"""Performs a node-level random split by adding :obj:`train_idx`,
    :obj:`val_idx` and :obj:`test_idx` attributes to the
    :class:`~graphlearn_torch.distributed.DistDataset` object. All nodes except 
    those in the validation and test sets will be used for training.
    
    Args:
      num_val (int or float): The number of validation samples.
        If float, it represents the ratio of samples to include in the
        validation set.
      num_test (int or float): The number of test samples in case
        of :obj:`"train_rest"` and :obj:`"random"` split. If float, it
        represents the ratio of samples to include in the test set.
    """

    if isinstance(self.node_labels, dict):
      train_idx = {}
      val_idx = {}
      test_idx = {}
  
      for node_type, _ in self.node_labels.items():
        indices = self.id_filter(self.node_pb[node_type], self.partition_idx)
        train_idx[node_type], val_idx[node_type], test_idx[node_type] = random_split(indices, num_val, num_test)
    else:
      indices = self.id_filter(self.node_pb, self.partition_idx)
      train_idx, val_idx, test_idx = random_split(indices, num_val, num_test)
    self.init_node_split((train_idx, val_idx, test_idx))

  def load_vineyard(
    self,
    vineyard_id: str,
    vineyard_socket: str,
    edges: List[EdgeType],
    edge_weights: Dict[EdgeType, str] = None,
    node_features: Dict[NodeType, List[str]] = None,
    edge_features: Dict[EdgeType, List[str]] = None,
    node_labels: Dict[NodeType, str] = None,
  ):
    super().load_vineyard(vineyard_id=vineyard_id, vineyard_socket=vineyard_socket, 
                          edges=edges, edge_weights=edge_weights, node_features=node_features, 
                          edge_features=edge_features, node_labels=node_labels)
    if isinstance(self.graph, dict):
      # hetero
      self._node_feat_pb = {}
      if node_features:
        for ntype, _ in self.node_features.items():
          if self.node_pb is not None:
            self._node_feat_pb[ntype] = self.node_pb[ntype]
          else:
            self._node_feat_pb[ntype] = None
    else:
      # homo
      if node_features:
        self._node_feat_pb = self.node_pb
    
    self.id_select = vineyard_utils.v6d_id_select
    self.id_filter = vineyard_utils.v6d_id_filter

  def share_ipc(self):
    super().share_ipc()
    self.node_pb = share_memory(self.node_pb)
    self.edge_pb = share_memory(self.edge_pb)
    self._node_feat_pb = share_memory(self._node_feat_pb)
    self._edge_feat_pb = share_memory(self._edge_feat_pb)
    ipc_hanlde = (
      self.num_partitions, self.partition_idx,
      self.graph, self.node_features, self.edge_features, self.node_labels,
      self.node_pb, self.edge_pb, self._node_feat_pb, self._edge_feat_pb, 
      self.edge_dir, self.graph_caching,
      (self.train_idx, self.val_idx, self.test_idx)
    )
    return ipc_hanlde

  @classmethod
  def from_ipc_handle(cls, ipc_handle):
    return cls(*ipc_handle)

  @property
  def node_feat_pb(self):
    if self._node_feat_pb is None:
      return self.node_pb
    return self._node_feat_pb

  @property
  def edge_feat_pb(self):
    if self._edge_feat_pb is None:
      return self.edge_pb
    return self._edge_feat_pb


def _cat_feature_cache(partition_idx, raw_feat_data, raw_feat_pb):
  r""" Cat a feature partition with its cached features.
  """
  if isinstance(raw_feat_data, dict):
    # heterogeneous.
    cache_ratio, feat_data, feat_id2idx, feat_pb = {}, {}, {}, {}
    for graph_type, raw_feat in raw_feat_data.items():
      cache_ratio[graph_type], feat_data[graph_type], \
      feat_id2idx[graph_type], feat_pb[graph_type] = \
        cat_feature_cache(partition_idx, raw_feat, raw_feat_pb[graph_type])
  else:
    # homogeneous.
    cache_ratio, feat_data, feat_id2idx, feat_pb = \
      cat_feature_cache(partition_idx, raw_feat_data, raw_feat_pb)
  return cache_ratio, feat_data, feat_id2idx, feat_pb


## Pickling Registration

def rebuild_dist_dataset(ipc_handle):
  ds = DistDataset.from_ipc_handle(ipc_handle)
  return ds

def reduce_dist_dataset(dataset: DistDataset):
  ipc_handle = dataset.share_ipc()
  return (rebuild_dist_dataset, (ipc_handle, ))

ForkingPickler.register(DistDataset, reduce_dist_dataset)

def random_split(
  indices: torch.Tensor,
  num_val: Union[float, int],
  num_test: Union[float, int],
):  
  num_total = indices.shape[0]
  num_val = round(num_total * num_val) if isinstance(num_val, float) else num_val
  num_test = round(num_total * num_test) if isinstance(num_test, float) else num_test
  perm = torch.randperm(num_total)
  val_idx = indices[perm[:num_val]].clone()
  test_idx = indices[perm[num_val:num_val + num_test]].clone()
  train_idx = indices[perm[num_val + num_test:]].clone()
  return train_idx, val_idx, test_idx