# 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. # ============================================================================== """ODPS table related distributed partitioner and dataset.""" import datetime from multiprocessing.reduction import ForkingPickler import numpy as np import torch import time from typing import Dict, Optional, Union try: import common_io except ImportError: pass from ..typing import ( NodeType, EdgeType, TensorDataType, ) from .dist_dataset import DistDataset, _cat_feature_cache from .dist_random_partitioner import DistRandomPartitioner class DistTableRandomPartitioner(DistRandomPartitioner): r""" A distributed random partitioner for parallel partitioning with large scale edge tables and node tables. Each distributed partitioner will process a slice of the full table, and partition them in parallel. After partitioning, each distributed partitioner will own a partitioned graph with its corresponding rank. Args: 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. Returns: int: Number of all partitions. int: The current partition idx. GraphPartitionData/HeteroGraphPartitionData: graph partition data. FeaturePartitionData/HeteroFeaturePartitionData: node feature partition data, optional. FeaturePartitionData/HeteroFeaturePartitionData: edge feature partition data, optional. PartitionBook/HeteroNodePartitionDict: node partition book. PartitionBook/HeteroEdgePartitionDict: edge partition book. """ def __init__( self, 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, ): super().__init__('', num_nodes, edge_index, edge_ids, node_feat, node_feat_ids, edge_feat, edge_feat_ids, num_parts, current_partition_idx, node_feat_dtype, edge_feat_dtype, edge_assign_strategy, chunk_size, master_addr, master_port, num_rpc_threads) 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. """ if 'hetero' == self.data_cls: node_pb_dict = {} node_feat_dict = {} for ntype in self.node_types: node_pb = self._partition_node(ntype) node_pb_dict[ntype] = node_pb current_node_feat_part = self._partition_node_feat(node_pb, ntype) if current_node_feat_part is not None: node_feat_dict[ntype] = current_node_feat_part edge_pb_dict = {} graph_dict = {} edge_feat_dict = {} for etype in self.edge_types: current_graph_part, edge_pb = self._partition_graph(node_pb_dict, etype) edge_pb_dict[etype] = edge_pb graph_dict[etype] = current_graph_part current_edge_feat_part = self._partition_edge_feat(edge_pb, etype) if current_edge_feat_part is not None: edge_feat_dict[etype] = current_edge_feat_part return ( self.num_parts, self.current_partition_idx, graph_dict, node_feat_dict, edge_feat_dict, node_pb_dict, edge_pb_dict ) else: node_pb = self._partition_node() node_feat = self._partition_node_feat(node_pb) graph, edge_pb = self._partition_graph(node_pb) edge_feat = self._partition_edge_feat(edge_pb) return ( self.num_parts, self.current_partition_idx, graph, node_feat, edge_feat, node_pb, edge_pb ) class DistTableDataset(DistDataset): """ Creates `DistDataset` from ODPS tables. Args: edge_tables: A dict({edge_type : odps_table}) denoting each bipartite graph input table of heterogeneous graph, where edge_type is a tuple of (src_type, edge_type, dst_type). node_tables: A dict({node_type(str) : odps_table}) denoting each input node table. num_nodes: Number of all graph nodes, should be a dict for hetero data. graph_mode: mode in graphlearn_torch's `Graph`, 'CPU', 'ZERO_COPY' or 'CUDA'. sort_func: function for feature reordering, return feature data(2D tenosr) and a map(1D tensor) from id to index. split_ratio: The proportion of data allocated to the GPU, between 0 and 1. device_group_list: A list of `DeviceGroup`. Each DeviceGroup must have the same size. A group of GPUs with peer-to-peer access to each other should be set in the same device group for high feature collection performance. directed: A Boolean value indicating whether the graph topology is directed. reader_threads: The number of threads of table reader. reader_capacity: The capacity of table reader. reader_batch_size: The number of records read at once. label: A CPU torch.Tensor(homo) or a Dict[NodeType, torch.Tensor](hetero) with the label data for graph nodes. device: The target cuda device rank to perform graph operations and feature lookups. 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``) 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 load( self, num_partitions=1, partition_idx=0, edge_tables=None, node_tables=None, num_nodes=0, graph_mode='ZERO_COPY', device_group_list=None, reader_threads=10, reader_capacity=10240, reader_batch_size=1024, label=None, device=None, feature_with_gpu=True, edge_assign_strategy='by_src', chunk_size=10000, master_addr=None, master_port=None, num_rpc_threads=16, ): assert isinstance(edge_tables, dict) assert isinstance(node_tables, dict) edge_index, eids, feature = {}, {}, {} edge_hetero = (len(edge_tables) > 1) node_hetero = (len(node_tables) > 1) print("Start Loading edge and node tables...") step = 0 start_time = time.time() for e_type, table in edge_tables.items(): edge_list = [] reader = common_io.table.TableReader(table, slice_id=partition_idx, slice_count=num_partitions, num_threads=reader_threads, capacity=reader_capacity) while True: try: data = reader.read(reader_batch_size, allow_smaller_final_batch=True) edge_list.extend(data) step += 1 except common_io.exception.OutOfRangeException: reader.close() break if step % 1000 == 0: print(f"{datetime.datetime.now()}: load " f"{step * reader_batch_size} edges.") rows = [e[0] for e in edge_list] cols = [e[1] for e in edge_list] eids_array = np.array([e[2] for e in edge_list], dtype=np.int64) edge_array = np.stack([np.array(rows, dtype=np.int64), np.array(cols, dtype=np.int64)]) if edge_hetero: edge_index[e_type] = eids_array eids[e_type] = eids else: edge_index = edge_array eids = eids_array del rows del cols del edge_list step = 0 for n_type, table in node_tables.items(): feature_list = [] reader = common_io.table.TableReader(table, slice_id=partition_idx, slice_count=num_partitions, num_threads=reader_threads, capacity=reader_capacity) while True: try: data = reader.read(reader_batch_size, allow_smaller_final_batch=True) feature_list.extend(data) step += 1 except common_io.exception.OutOfRangeException: reader.close() break if step % 1000 == 0: print(f"{datetime.datetime.now()}: load " f"{step * reader_batch_size} nodes.") ids = torch.tensor([feat[0] for feat in feature_list], dtype=torch.long) if isinstance(feature_list[0][1], bytes): float_feat= [ list(map(float, feat[1].decode().split(':'))) for feat in feature_list ] else: float_feat= [ list(map(float, feat[1].split(':'))) for feat in feature_list ] if node_hetero: feature[n_type] = torch.tensor(float_feat) else: feature = torch.tensor(float_feat) del float_feat del feature_list load_time = (time.time() - start_time) / 60 print(f'Loading table completed in {load_time:.2f} minutes.') print("Start partitioning graph and feature...") p_start = time.time() dist_partitioner = DistTableRandomPartitioner( num_nodes, edge_index=edge_index, edge_ids=eids, node_feat=feature, node_feat_ids=ids, num_parts=num_partitions, current_partition_idx=partition_idx, edge_assign_strategy=edge_assign_strategy, chunk_size=chunk_size, master_addr=master_addr, master_port=master_port, num_rpc_threads=num_rpc_threads) ( self.num_partitions, self.partition_idx, graph_data, node_feat_data, edge_feat_data, self.node_pb, self.edge_pb ) = dist_partitioner.partition() part_time = (time.time() - p_start) / 60 print(f'Partitioning completed in {part_time:.2f} minutes.') # init graph if isinstance(graph_data, dict): # heterogeneous. edge_index, edge_ids = {}, {} for k, v in graph_data.items(): edge_index[k] = v.edge_index edge_ids[k] = v.eids else: # homogeneous. edge_index = graph_data.edge_index edge_ids = graph_data.eids self.init_graph(edge_index, edge_ids, layout='COO', graph_mode=graph_mode, device=device) # load node feature 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 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 self.init_node_labels(label) ## Pickling Registration def rebuild_dist_table_dataset(ipc_handle): ds = DistTableDataset.from_ipc_handle(ipc_handle) return ds def reduce_dist_table_dataset(dataset: DistTableDataset): ipc_handle = dataset.share_ipc() return (rebuild_dist_table_dataset, (ipc_handle, )) ForkingPickler.register(DistTableDataset, reduce_dist_table_dataset)