graphlearn_torch/python/distributed/dist_random_partitioner.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, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import threading
from typing import Any, Dict, List, Optional, Tuple, Union

import torch

from ..partition import (
  save_meta, save_node_pb, save_edge_pb,
  save_graph_partition, save_feature_partition,
  PartitionBook
)
from ..typing import (
  NodeType, EdgeType, TensorDataType,
  GraphPartitionData, FeaturePartitionData
)
from ..utils import convert_to_tensor, ensure_dir, index_select

from .dist_context import get_context, init_worker_group
from .rpc import (
  init_rpc, rpc_is_initialized, all_gather, barrier,
  get_rpc_current_group_worker_names,
  rpc_request_async, rpc_register, RpcCalleeBase
)


class RpcUpdatePartitionValueCallee(RpcCalleeBase):
  def __init__(self, dist_partition_mgr):
    super().__init__()
    self.dist_partition_mgr = dist_partition_mgr

  def call(self, *args, **kwargs):
    self.dist_partition_mgr._update_part_val(*args, **kwargs)
    return None


class RpcUpdatePartitionBookCallee(RpcCalleeBase):
  def __init__(self, dist_partition_mgr):
    super().__init__()
    self.dist_partition_mgr = dist_partition_mgr

  def call(self, *args, **kwargs):
    self.dist_partition_mgr._update_pb(*args, **kwargs)
    return None


class DistPartitionManager(object):
  r""" A state manager for distributed partitioning.
  """
  def __init__(self, total_val_size: int = 1, generate_pb: bool = True):
    assert rpc_is_initialized() is True
    self.num_parts = get_context().world_size
    self.cur_pidx = get_context().rank

    self._lock = threading.RLock()
    self._worker_names = get_rpc_current_group_worker_names()

    self.reset(total_val_size, generate_pb)

    val_update_callee = RpcUpdatePartitionValueCallee(self)
    self._val_update_callee_id = rpc_register(val_update_callee)
    pb_update_callee = RpcUpdatePartitionBookCallee(self)
    self._pb_update_callee_id = rpc_register(pb_update_callee)

  def reset(self, total_val_size: int, generate_pb: bool = True):
    with self._lock:
      self.generate_pb = generate_pb
      self.cur_part_val_list = []
      if self.generate_pb:
        self.partition_book = torch.zeros(total_val_size, dtype=torch.int64)
      else:
        self.partition_book = None

  def process(self, res_list: List[Tuple[Any, torch.Tensor]]):
    r""" Process partitioned results of the current corresponded distributed
    partitioner and synchronize with others.

    Args:
      res_list: The result list of value and ids for each partition.
    """
    assert len(res_list) == self.num_parts
    futs = []
    for pidx, (val, val_idx) in enumerate(res_list):
      if pidx == self.cur_pidx:
        self._update_part_val(val, pidx)
      else:
        futs.append(rpc_request_async(self._worker_names[pidx],
                                      self._val_update_callee_id,
                                      args=(val, pidx)))
      if self.generate_pb:
        futs.extend(self._broadcast_pb(val_idx, pidx))
    _ = torch.futures.wait_all(futs)

  def _broadcast_pb(self, val_idx: torch.Tensor, target_pidx: int):
    pb_update_futs = []
    for pidx in range(self.num_parts):
      if pidx == self.cur_pidx:
        self._update_pb(val_idx, target_pidx)
      else:
        pb_update_futs.append(rpc_request_async(self._worker_names[pidx],
                                                self._pb_update_callee_id,
                                                args=(val_idx, target_pidx)))
    return pb_update_futs

  def _update_part_val(self, val, target_pidx: int):
    assert target_pidx == self.cur_pidx
    with self._lock:
      if val is not None:
        self.cur_part_val_list.append(val)

  def _update_pb(self, val_idx: torch.Tensor, target_pidx: int):
    with self._lock:
      self.partition_book[val_idx] = target_pidx


class DistRandomPartitioner(object):
  r""" A distributed random partitioner for parallel partitioning with large
  scale graph and features.

  Each distributed partitioner will process a part of the full graph and
  feature data, and partition them in parallel. A distributed partitioner's
  rank is corresponding to a partition index, and the number of all distributed
  partitioners must be same with the number of output partitions. During
  partitioning, the partitioned results will be sent to other distributed
  partitioners according to their ranks. After partitioning, each distributed
  partitioner will own a partitioned graph with its corresponding rank and
  further save the partitioned results into the local output directory.

  Args:
    output_dir: The output root directory on local machine for partitioned
      results.
    num_nodes: Number of all graph nodes, should be a dict for hetero data.
    edge_index: A part of the edge index data of graph edges, should be a dict
      for hetero data.
    edge_ids: The edge ids of the input ``edge_index``.
    node_feat: A part of the node feature data, should be a dict for hetero data.
    node_feat_ids: The node ids corresponding to the input ``node_feat``.
    edge_feat: The edge feature data, should be a dict for hetero data.
    edge_feat_ids: The edge ids corresponding to the input ``edge_feat``.
    num_parts: The number of all partitions. If not provided, the value of
      ``graphlearn_torch.distributed.get_context().world_size`` will be used.
    current_partition_idx: The partition index corresponding to the current
      distributed partitioner. If not provided, the value of
      ``graphlearn_torch.distributed.get_context().rank`` will be used.
    node_feat_dtype: The data type of node features.
    edge_feat_dtype: The data type of edge features.
    edge_assign_strategy: The assignment strategy when partitioning edges,
      should be 'by_src' or 'by_dst'.
    chunk_size: The chunk size for partitioning.
    master_addr: The master TCP address for RPC connection between all
      distributed partitioners.
    master_port: The master TCP port for RPC connection between all
      distributed partitioners.
    num_rpc_threads: The number of RPC worker threads to use.
  """
  def __init__(
    self,
    output_dir: str,
    num_nodes: Union[int, Dict[NodeType, int]],
    edge_index: Union[TensorDataType, Dict[EdgeType, TensorDataType]],
    edge_ids: Union[TensorDataType, Dict[EdgeType, TensorDataType]],
    node_feat: Optional[Union[TensorDataType, Dict[NodeType, TensorDataType]]] = None,
    node_feat_ids: Optional[Union[TensorDataType, Dict[NodeType, TensorDataType]]] = None,
    edge_feat: Optional[Union[TensorDataType, Dict[EdgeType, TensorDataType]]] = None,
    edge_feat_ids: Optional[Union[TensorDataType, Dict[EdgeType, TensorDataType]]] = None,
    num_parts: Optional[int] = None,
    current_partition_idx: Optional[int] = None,
    node_feat_dtype: torch.dtype = torch.float32,
    edge_feat_dtype: torch.dtype = torch.float32,
    edge_assign_strategy: str = 'by_src',
    chunk_size: int = 10000,
    master_addr: Optional[str] = None,
    master_port: Optional[str] = None,
    num_rpc_threads: int = 16,
  ):
    self.output_dir = output_dir

    if get_context() is not None:
      if num_parts is not None:
        assert get_context().world_size == num_parts
      if current_partition_idx is not None:
        assert get_context().rank == current_partition_idx
    else:
      assert num_parts is not None
      assert current_partition_idx is not None
      init_worker_group(
        world_size=num_parts,
        rank=current_partition_idx,
        group_name='distributed_random_partitoner'
      )
    self.num_parts = get_context().world_size
    self.current_partition_idx = get_context().rank

    if rpc_is_initialized() is not True:
      if master_addr is None:
        master_addr = os.environ['MASTER_ADDR']
      if master_port is None:
        master_port = int(os.environ['MASTER_PORT'])
      init_rpc(master_addr, master_port, num_rpc_threads)

    self.num_nodes = num_nodes
    self.edge_index = convert_to_tensor(edge_index, dtype=torch.int64)
    self.edge_ids = convert_to_tensor(edge_ids, dtype=torch.int64)

    self.node_feat = convert_to_tensor(node_feat, dtype=node_feat_dtype)
    self.node_feat_ids = convert_to_tensor(node_feat_ids, dtype=torch.int64)
    if self.node_feat is not None:
      assert self.node_feat_ids is not None

    self.edge_feat = convert_to_tensor(edge_feat, dtype=edge_feat_dtype)
    self.edge_feat_ids = convert_to_tensor(edge_feat_ids, dtype=torch.int64)
    if self.edge_feat is not None:
      assert self.edge_feat_ids is not None

    if isinstance(self.num_nodes, dict):
      assert isinstance(self.edge_index, dict)
      assert isinstance(self.edge_ids, dict)
      assert isinstance(self.node_feat, dict) or self.node_feat is None
      assert isinstance(self.node_feat_ids, dict) or self.node_feat_ids is None
      assert isinstance(self.edge_feat, dict) or self.edge_feat is None
      assert isinstance(self.edge_feat_ids, dict) or self.edge_feat_ids is None
      self.data_cls = 'hetero'
      self.node_types = sorted(list(self.num_nodes.keys()))
      self.edge_types = sorted(list(self.edge_index.keys()))
      self.num_local_edges = {}
      self.num_edges = {}
      for etype, index in self.edge_index.items():
        self.num_local_edges[etype] = len(index[0])
        self.num_edges[etype] = sum(all_gather(len(index[0])).values())
    else:
      self.data_cls = 'homo'
      self.node_types = None
      self.edge_types = None
      self.num_local_edges = len(self.edge_index[0])
      self.num_edges = sum(all_gather(len(self.edge_index[0])).values())

    self.edge_assign_strategy = edge_assign_strategy.lower()
    assert self.edge_assign_strategy in ['by_src', 'by_dst']
    self.chunk_size = chunk_size

    self._partition_mgr = DistPartitionManager()

  def _partition_by_chunk(
    self,
    val: Any,
    val_idx: torch.Tensor,
    partition_fn,
    total_val_size: int,
    generate_pb = True
  ):
    r""" Partition generic values and sync with all other partitoners.
    """
    val_num = len(val_idx)
    chunk_num = (val_num + self.chunk_size - 1) // self.chunk_size
    chunk_start_pos = 0
    self._partition_mgr.reset(total_val_size, generate_pb)
    barrier()
    for _ in range(chunk_num):
      chunk_end_pos = min(val_num, chunk_start_pos + self.chunk_size)
      current_chunk_size = chunk_end_pos - chunk_start_pos
      chunk_idx = torch.arange(current_chunk_size, dtype=torch.long)
      chunk_val = index_select(val, index=(chunk_start_pos, chunk_end_pos))
      chunk_val_idx = val_idx[chunk_start_pos:chunk_end_pos]
      chunk_partition_idx = partition_fn(
        chunk_val_idx, (chunk_start_pos, chunk_end_pos))
      chunk_res = []
      for pidx in range(self.num_parts):
        mask = (chunk_partition_idx == pidx)
        idx = torch.masked_select(chunk_idx, mask)
        chunk_res.append((index_select(chunk_val, idx), chunk_val_idx[idx]))
      self._partition_mgr.process(chunk_res)
      chunk_start_pos += current_chunk_size
    barrier()
    return (
      self._partition_mgr.cur_part_val_list,
      self._partition_mgr.partition_book
    )

  def _partition_node(
    self,
    ntype: Optional[NodeType] = None
  ) -> PartitionBook:
    r""" Partition graph nodes of a specify node type in parallel.

    Args:
      ntype (str): The type for input nodes, must be provided for heterogeneous
        graph. (default: ``None``)

    Returns:
      PartitionBook: The partition book of graph nodes.
    """
    if 'hetero' == self.data_cls:
      assert ntype is not None
      node_num = self.num_nodes[ntype]
    else:
      node_num = self.num_nodes
    per_node_num = node_num // self.num_parts
    local_node_start = per_node_num * self.current_partition_idx
    local_node_end = min(
      node_num,
      per_node_num * (self.current_partition_idx + 1)
    )
    local_node_ids = torch.arange(
      local_node_start, local_node_end, dtype=torch.int64
    )

    def _node_pfn(n_ids, _):
      partition_idx = n_ids % self.num_parts
      rand_order = torch.randperm(len(n_ids))
      return partition_idx[rand_order]

    _, node_pb = self._partition_by_chunk(
      val=None,
      val_idx=local_node_ids,
      partition_fn=_node_pfn,
      total_val_size=node_num,
      generate_pb=True
    )
    return node_pb

  def _partition_graph(
    self,
    node_pbs: Union[PartitionBook, Dict[NodeType, PartitionBook]],
    etype: Optional[EdgeType] = None
  ) -> Tuple[GraphPartitionData, PartitionBook]:
    r""" Partition graph topology of a specify edge type in parallel.

    Args:
      node_pbs (PartitionBook or Dict[NodeType, PartitionBook]): The
        partition books of all graph nodes.
      etype (Tuple[str, str, str]): The type for input edges, must be provided
        for heterogeneous graph. (default: ``None``)

    Returns:
      GraphPartitionData: The graph data of the current partition.
      PartitionBook: The partition book of graph edges.
    """
    if 'hetero' == self.data_cls:
      assert isinstance(node_pbs, dict)
      assert etype is not None
      src_ntype, _, dst_ntype = etype

      edge_index = self.edge_index[etype]
      rows, cols = edge_index[0], edge_index[1]
      eids = self.edge_ids[etype]
      edge_num = self.num_edges[etype]

      if 'by_src' == self.edge_assign_strategy:
        target_node_pb = node_pbs[src_ntype]
        target_indices = rows
      else:
        target_node_pb = node_pbs[dst_ntype]
        target_indices = cols
    else:
      edge_index = self.edge_index
      rows, cols = edge_index[0], edge_index[1]
      eids = self.edge_ids
      edge_num = self.num_edges

      target_node_pb = node_pbs
      target_indices = rows if 'by_src' == self.edge_assign_strategy else cols

    def _edge_pfn(_, chunk_range):
      chunk_target_indices = index_select(target_indices, chunk_range)
      return target_node_pb[chunk_target_indices]

    res_list, edge_pb = self._partition_by_chunk(
      val=(rows, cols, eids),
      val_idx=eids,
      partition_fn=_edge_pfn,
      total_val_size=edge_num,
      generate_pb=True
    )
    current_graph_part = GraphPartitionData(
      edge_index=(
        torch.cat([r[0] for r in res_list]),
        torch.cat([r[1] for r in res_list]),
      ),
      eids=torch.cat([r[2] for r in res_list])
    )
    return current_graph_part, edge_pb

  def _partition_node_feat(
    self,
    node_pb: PartitionBook,
    ntype: Optional[NodeType] = None,
  ) -> Optional[FeaturePartitionData]:
    r""" Split node features in parallel by the partitioned node results,
    and return the current partition of node features.
    """
    if self.node_feat is None:
      return None

    if 'hetero' == self.data_cls:
      assert ntype is not None
      node_num = self.num_nodes[ntype]
      node_feat = self.node_feat[ntype]
      node_feat_ids = self.node_feat_ids[ntype]
    else:
      node_num = self.num_nodes
      node_feat = self.node_feat
      node_feat_ids = self.node_feat_ids

    def _node_feat_pfn(nfeat_ids, _):
      return node_pb[nfeat_ids]

    res_list, _ = self._partition_by_chunk(
      val=(node_feat, node_feat_ids),
      val_idx=node_feat_ids,
      partition_fn=_node_feat_pfn,
      total_val_size=node_num,
      generate_pb=False
    )
    return FeaturePartitionData(
      feats=torch.cat([r[0] for r in res_list]),
      ids=torch.cat([r[1] for r in res_list]),
      cache_feats=None,
      cache_ids=None
    )

  def _partition_edge_feat(
    self,
    edge_pb: PartitionBook,
    etype: Optional[EdgeType] = None,
  ) -> Optional[FeaturePartitionData]:
    r""" Split edge features in parallel by the partitioned edge results,
    and return the current partition of edge features.
    """
    if self.edge_feat is None:
      return None

    if 'hetero' == self.data_cls:
      assert etype is not None
      edge_num = self.num_edges[etype]
      edge_feat = self.edge_feat[etype]
      edge_feat_ids = self.edge_feat_ids[etype]
    else:
      edge_num = self.num_edges
      edge_feat = self.edge_feat
      edge_feat_ids = self.edge_feat_ids

    def _edge_feat_pfn(efeat_ids, _):
      return edge_pb[efeat_ids]

    res_list, _ = self._partition_by_chunk(
      val=(edge_feat, edge_feat_ids),
      val_idx=edge_feat_ids,
      partition_fn=_edge_feat_pfn,
      total_val_size=edge_num,
      generate_pb=False
    )
    return FeaturePartitionData(
      feats=torch.cat([r[0] for r in res_list]),
      ids=torch.cat([r[1] for r in res_list]),
      cache_feats=None,
      cache_ids=None
    )

  def partition(self):
    r""" Partition graph and feature data into different parts along with all
    other distributed partitioners, save the result of the current partition
    index into output directory.
    """
    ensure_dir(self.output_dir)
    if 'hetero' == self.data_cls:
      node_pb_dict = {}
      for ntype in self.node_types:
        node_pb = self._partition_node(ntype)
        node_pb_dict[ntype] = node_pb
        save_node_pb(self.output_dir, node_pb, ntype)

        current_node_feat_part = self._partition_node_feat(node_pb, ntype)
        if current_node_feat_part is not None:
          save_feature_partition(
            self.output_dir, self.current_partition_idx, current_node_feat_part,
            group='node_feat', graph_type=ntype
          )
        del current_node_feat_part

      for etype in self.edge_types:
        current_graph_part, edge_pb = self._partition_graph(node_pb_dict, etype)
        save_edge_pb(self.output_dir, edge_pb, etype)
        save_graph_partition(
          self.output_dir, self.current_partition_idx, current_graph_part, etype
        )
        del current_graph_part

        current_edge_feat_part = self._partition_edge_feat(edge_pb, etype)
        if current_edge_feat_part is not None:
          save_feature_partition(
            self.output_dir, self.current_partition_idx, current_edge_feat_part,
            group='edge_feat', graph_type=etype
          )
        del current_edge_feat_part

    else:
      node_pb = self._partition_node()
      save_node_pb(self.output_dir, node_pb)

      current_node_feat_part = self._partition_node_feat(node_pb)
      if current_node_feat_part is not None:
        save_feature_partition(
          self.output_dir, self.current_partition_idx,
          current_node_feat_part, group='node_feat'
        )
      del current_node_feat_part

      current_graph_part, edge_pb = self._partition_graph(node_pb)
      save_edge_pb(self.output_dir, edge_pb)
      save_graph_partition(
        self.output_dir, self.current_partition_idx, current_graph_part
      )
      del current_graph_part

      current_edge_feat_part = self._partition_edge_feat(edge_pb)
      if current_edge_feat_part is not None:
        save_feature_partition(
          self.output_dir, self.current_partition_idx,
          current_edge_feat_part, group='edge_feat'
        )
      del current_edge_feat_part

    # save meta.
    save_meta(self.output_dir, self.num_parts, self.data_cls,
              self.node_types, self.edge_types)