"""A set of graph services of getting subgraphs from DistGraph"""importosfromcollectionsimportnamedtupleimportnumpyasnpimporttorchfrom..importbackendasF,graphboltasgbfrom..baseimportEID,ETYPE,NIDfrom..convertimportgraph,heterographfrom..samplingimport(sample_etype_neighborsaslocal_sample_etype_neighbors,sample_neighborsaslocal_sample_neighbors,)from..subgraphimportin_subgraphaslocal_in_subgraphfrom..utilsimporttoindexfrom.constantsimportDGL2GB_EID,GB_DST_IDfrom.rpcimport(recv_responses,register_service,Request,Response,send_requests_to_machine,)__all__=["sample_neighbors","sample_etype_neighbors","in_subgraph","find_edges",]SAMPLING_SERVICE_ID=6657INSUBGRAPH_SERVICE_ID=6658EDGES_SERVICE_ID=6659OUTDEGREE_SERVICE_ID=6660INDEGREE_SERVICE_ID=6661ETYPE_SAMPLING_SERVICE_ID=6662classSubgraphResponse(Response):"""The response for sampling and in_subgraph"""def__init__(self,global_src,global_dst,*,global_eids=None,etype_ids=None):self.global_src=global_srcself.global_dst=global_dstself.global_eids=global_eidsself.etype_ids=etype_idsdef__setstate__(self,state):(self.global_src,self.global_dst,self.global_eids,self.etype_ids,)=statedef__getstate__(self):return(self.global_src,self.global_dst,self.global_eids,self.etype_ids,)classFindEdgeResponse(Response):"""The response for sampling and in_subgraph"""def__init__(self,global_src,global_dst,order_id):self.global_src=global_srcself.global_dst=global_dstself.order_id=order_iddef__setstate__(self,state):self.global_src,self.global_dst,self.order_id=statedef__getstate__(self):returnself.global_src,self.global_dst,self.order_iddef_sample_neighbors_graphbolt(g,gpb,nodes,fanout,edge_dir="in",prob=None,exclude_edges=None,replace=False,):"""Sample from local partition via graphbolt. The input nodes use global IDs. We need to map the global node IDs to local node IDs, perform sampling and map the sampled results to the global IDs space again. The sampled results are stored in three vectors that store source nodes, destination nodes, etype IDs and edge IDs. Parameters ---------- g : FusedCSCSamplingGraph The local partition. gpb : GraphPartitionBook The graph partition book. nodes : tensor The nodes to sample neighbors from. fanout : tensor or int The number of edges to be sampled for each node. edge_dir : str, optional Determines whether to sample inbound or outbound edges. prob : tensor, optional The probability associated with each neighboring edge of a node. exclude_edges : tensor, optional The edges to exclude when sampling. replace : bool, optional If True, sample with replacement. Returns ------- tensor The source node ID array. tensor The destination node ID array. tensor The edge ID array. tensor The edge type ID array. """assert(edge_dir=="in"),f"GraphBolt only supports inbound edge sampling but got {edge_dir}."assertexclude_edgesisNone,"GraphBolt does not support excluding edges."# 1. Map global node IDs to local node IDs.nodes=gpb.nid2localnid(nodes,gpb.partid)# Local partition may be saved in torch.int32 even though the global graph# is in torch.int64.nodes=nodes.to(dtype=g.indices.dtype)# 2. Perform sampling.probs_or_mask=NoneifprobisnotNone:probs_or_mask=g.edge_attributes[prob]# Sanity checks.assertisinstance(g,gb.FusedCSCSamplingGraph),"Expect a FusedCSCSamplingGraph."assertisinstance(nodes,torch.Tensor),"Expect a tensor of nodes."ifisinstance(fanout,int):fanout=torch.LongTensor([fanout])assertisinstance(fanout,torch.Tensor),"Expect a tensor of fanout."subgraph=g._sample_neighbors(nodes,None,fanout,replace=replace,probs_or_mask=probs_or_mask,)# 3. Map local node IDs to global node IDs.local_src=subgraph.indiceslocal_dst=gb.expand_indptr(subgraph.indptr,dtype=local_src.dtype,node_ids=subgraph.original_column_node_ids,output_size=local_src.shape[0],)global_nid_mapping=g.node_attributes[NID]global_src=global_nid_mapping[local_src]global_dst=global_nid_mapping[local_dst]global_eids=Noneifg.edge_attributesisnotNoneandEIDing.edge_attributes:global_eids=g.edge_attributes[EID][subgraph.original_edge_ids]returnLocalSampledGraph(global_src,global_dst,global_eids,subgraph.type_per_edge)def_sample_neighbors_dgl(local_g,partition_book,seed_nodes,fan_out,edge_dir="in",prob=None,exclude_edges=None,replace=False,):"""Sample from local partition. The input nodes use global IDs. We need to map the global node IDs to local node IDs, perform sampling and map the sampled results to the global IDs space again. The sampled results are stored in three vectors that store source nodes, destination nodes and edge IDs. """local_ids=partition_book.nid2localnid(seed_nodes,partition_book.partid)local_ids=F.astype(local_ids,local_g.idtype)# local_ids = self.seed_nodessampled_graph=local_sample_neighbors(local_g,local_ids,fan_out,edge_dir=edge_dir,prob=prob,exclude_edges=exclude_edges,replace=replace,_dist_training=True,)global_nid_mapping=local_g.ndata[NID]src,dst=sampled_graph.edges()global_src,global_dst=F.gather_row(global_nid_mapping,src),F.gather_row(global_nid_mapping,dst)global_eids=F.gather_row(local_g.edata[EID],sampled_graph.edata[EID])returnLocalSampledGraph(global_src,global_dst,global_eids)def_sample_neighbors(use_graphbolt,*args,**kwargs):"""Wrapper for sampling neighbors. The actual sampling function depends on whether to use GraphBolt. Parameters ---------- use_graphbolt : bool Whether to use GraphBolt for sampling. args : list The arguments for the sampling function. kwargs : dict The keyword arguments for the sampling function. Returns ------- tensor The source node ID array. tensor The destination node ID array. tensor The edge ID array. tensor The edge type ID array. """func=(_sample_neighbors_graphboltifuse_graphboltelse_sample_neighbors_dgl)returnfunc(*args,**kwargs)def_sample_etype_neighbors_dgl(local_g,partition_book,seed_nodes,fan_out,edge_dir="in",prob=None,exclude_edges=None,replace=False,etype_offset=None,etype_sorted=False,):"""Sample from local partition. The input nodes use global IDs. We need to map the global node IDs to local node IDs, perform sampling and map the sampled results to the global IDs space again. The sampled results are stored in three vectors that store source nodes, destination nodes and edge IDs. """assertetype_offsetisnotNone,"The etype offset is not provided."local_ids=partition_book.nid2localnid(seed_nodes,partition_book.partid)local_ids=F.astype(local_ids,local_g.idtype)sampled_graph=local_sample_etype_neighbors(local_g,local_ids,etype_offset,fan_out,edge_dir=edge_dir,prob=prob,exclude_edges=exclude_edges,replace=replace,etype_sorted=etype_sorted,_dist_training=True,)global_nid_mapping=local_g.ndata[NID]src,dst=sampled_graph.edges()global_src,global_dst=F.gather_row(global_nid_mapping,src),F.gather_row(global_nid_mapping,dst)global_eids=F.gather_row(local_g.edata[EID],sampled_graph.edata[EID])returnLocalSampledGraph(global_src,global_dst,global_eids)def_sample_etype_neighbors(use_graphbolt,*args,**kwargs):"""Wrapper for sampling etype neighbors. The actual sampling function depends on whether to use GraphBolt. Parameters ---------- use_graphbolt : bool Whether to use GraphBolt for sampling. args : list The arguments for the sampling function. kwargs : dict The keyword arguments for the sampling function. Returns ------- tensor The source node ID array. tensor The destination node ID array. tensor The edge ID array. tensor The edge type ID array. """func=(_sample_neighbors_graphboltifuse_graphboltelse_sample_etype_neighbors_dgl)ifuse_graphbolt:# GraphBolt does not require `etype_offset` and `etype_sorted`.kwargs.pop("etype_offset",None)kwargs.pop("etype_sorted",None)returnfunc(*args,**kwargs)def_find_edges(local_g,partition_book,seed_edges):"""Given an edge ID array, return the source and destination node ID array ``s`` and ``d`` in the local partition. """local_eids=partition_book.eid2localeid(seed_edges,partition_book.partid)ifisinstance(local_g,gb.FusedCSCSamplingGraph):# When converting from DGLGraph to FusedCSCSamplingGraph, the edge IDs# are re-ordered. In order to find the correct node pairs, we need to# map the DGL edge IDs back to GraphBolt edge IDs.if(DGL2GB_EIDnotinlocal_g.edge_attributesorGB_DST_IDnotinlocal_g.edge_attributes):raiseValueError("The edge attributes DGL2GB_EID and GB_DST_ID are not found. ""Please make sure `coo` format is available when generating ""partitions in GraphBolt format.")local_eids=local_g.edge_attributes[DGL2GB_EID][local_eids]local_src=local_g.indices[local_eids]local_dst=local_g.edge_attributes[GB_DST_ID][local_eids]global_nid_mapping=local_g.node_attributes[NID]else:local_eids=F.astype(local_eids,local_g.idtype)local_src,local_dst=local_g.find_edges(local_eids)global_nid_mapping=local_g.ndata[NID]global_src=global_nid_mapping[local_src]global_dst=global_nid_mapping[local_dst]returnglobal_src,global_dstdef_in_degrees(local_g,partition_book,n):"""Get in-degree of the nodes in the local partition."""local_nids=partition_book.nid2localnid(n,partition_book.partid)local_nids=F.astype(local_nids,local_g.idtype)returnlocal_g.in_degrees(local_nids)def_out_degrees(local_g,partition_book,n):"""Get out-degree of the nodes in the local partition."""local_nids=partition_book.nid2localnid(n,partition_book.partid)local_nids=F.astype(local_nids,local_g.idtype)returnlocal_g.out_degrees(local_nids)def_in_subgraph(local_g,partition_book,seed_nodes):"""Get in subgraph from local partition. The input nodes use global IDs. We need to map the global node IDs to local node IDs, get in-subgraph and map the sampled results to the global IDs space again. The results are stored in three vectors that store source nodes, destination nodes and edge IDs. """local_ids=partition_book.nid2localnid(seed_nodes,partition_book.partid)local_ids=F.astype(local_ids,local_g.idtype)# local_ids = self.seed_nodessampled_graph=local_in_subgraph(local_g,local_ids)global_nid_mapping=local_g.ndata[NID]src,dst=sampled_graph.edges()global_src,global_dst=global_nid_mapping[src],global_nid_mapping[dst]global_eids=F.gather_row(local_g.edata[EID],sampled_graph.edata[EID])returnLocalSampledGraph(global_src,global_dst,global_eids)# --- NOTE 1 ---# (BarclayII)# If the sampling algorithm needs node and edge data, ideally the# algorithm should query the underlying feature storage to get what it# just needs to complete the job. For instance, with# sample_etype_neighbors, we only need the probability of the seed nodes'# neighbors.## However, right now we are reusing the existing subgraph sampling# interfaces of DGLGraph (i.e. single machine solution), which needs# the data of *all* the nodes/edges. Going distributed, we now need# the node/edge data of the *entire* local graph partition.## If the sampling algorithm only use edge data, the current design works# because the local graph partition contains all the in-edges of the# assigned nodes as well as the data. This is the case for# sample_etype_neighbors.## However, if the sampling algorithm requires data of the neighbor nodes# (e.g. sample_neighbors_biased which performs biased sampling based on the# type of the neighbor nodes), the current design will fail because the# neighbor nodes (hence the data) may not belong to the current partition.# This is a limitation of the current DistDGL design. We should improve it# later.classSamplingRequest(Request):"""Sampling Request"""def__init__(self,nodes,fan_out,edge_dir="in",prob=None,exclude_edges=None,replace=False,use_graphbolt=False,):self.seed_nodes=nodesself.edge_dir=edge_dirself.prob=probself.exclude_edges=exclude_edgesself.replace=replaceself.fan_out=fan_outself.use_graphbolt=use_graphboltdef__setstate__(self,state):(self.seed_nodes,self.edge_dir,self.prob,self.exclude_edges,self.replace,self.fan_out,self.use_graphbolt,)=statedef__getstate__(self):return(self.seed_nodes,self.edge_dir,self.prob,self.exclude_edges,self.replace,self.fan_out,self.use_graphbolt,)defprocess_request(self,server_state):local_g=server_state.graphpartition_book=server_state.partition_bookkv_store=server_state.kv_storeifself.probisnotNoneand(notself.use_graphbolt):prob=[kv_store.data_store[self.prob]]else:prob=self.probres=_sample_neighbors(self.use_graphbolt,local_g,partition_book,self.seed_nodes,self.fan_out,edge_dir=self.edge_dir,prob=prob,exclude_edges=self.exclude_edges,replace=self.replace,)returnSubgraphResponse(res.global_src,res.global_dst,global_eids=res.global_eids,etype_ids=res.etype_ids,)classSamplingRequestEtype(Request):"""Sampling Request"""def__init__(self,nodes,fan_out,edge_dir="in",prob=None,exclude_edges=None,replace=False,etype_sorted=True,use_graphbolt=False,):self.seed_nodes=nodesself.edge_dir=edge_dirself.prob=probself.exclude_edges=exclude_edgesself.replace=replaceself.fan_out=fan_outself.etype_sorted=etype_sortedself.use_graphbolt=use_graphboltdef__setstate__(self,state):(self.seed_nodes,self.edge_dir,self.prob,self.exclude_edges,self.replace,self.fan_out,self.etype_sorted,self.use_graphbolt,)=statedef__getstate__(self):return(self.seed_nodes,self.edge_dir,self.prob,self.exclude_edges,self.replace,self.fan_out,self.etype_sorted,self.use_graphbolt,)defprocess_request(self,server_state):local_g=server_state.graphpartition_book=server_state.partition_bookkv_store=server_state.kv_storeetype_offset=partition_book.local_etype_offset# See NOTE 1ifself.probisnotNoneand(notself.use_graphbolt):probs=[kv_store.data_store[key]ifkey!=""elseNoneforkeyinself.prob]else:probs=self.probres=_sample_etype_neighbors(self.use_graphbolt,local_g,partition_book,self.seed_nodes,self.fan_out,edge_dir=self.edge_dir,prob=probs,exclude_edges=self.exclude_edges,replace=self.replace,etype_offset=etype_offset,etype_sorted=self.etype_sorted,)returnSubgraphResponse(res.global_src,res.global_dst,global_eids=res.global_eids,etype_ids=res.etype_ids,)classEdgesRequest(Request):"""Edges Request"""def__init__(self,edge_ids,order_id):self.edge_ids=edge_idsself.order_id=order_iddef__setstate__(self,state):self.edge_ids,self.order_id=statedef__getstate__(self):returnself.edge_ids,self.order_iddefprocess_request(self,server_state):local_g=server_state.graphpartition_book=server_state.partition_bookglobal_src,global_dst=_find_edges(local_g,partition_book,self.edge_ids)returnFindEdgeResponse(global_src,global_dst,self.order_id)classInDegreeRequest(Request):"""In-degree Request"""def__init__(self,n,order_id):self.n=nself.order_id=order_iddef__setstate__(self,state):self.n,self.order_id=statedef__getstate__(self):returnself.n,self.order_iddefprocess_request(self,server_state):local_g=server_state.graphpartition_book=server_state.partition_bookdeg=_in_degrees(local_g,partition_book,self.n)returnInDegreeResponse(deg,self.order_id)classInDegreeResponse(Response):"""The response for in-degree"""def__init__(self,deg,order_id):self.val=degself.order_id=order_iddef__setstate__(self,state):self.val,self.order_id=statedef__getstate__(self):returnself.val,self.order_idclassOutDegreeRequest(Request):"""Out-degree Request"""def__init__(self,n,order_id):self.n=nself.order_id=order_iddef__setstate__(self,state):self.n,self.order_id=statedef__getstate__(self):returnself.n,self.order_iddefprocess_request(self,server_state):local_g=server_state.graphpartition_book=server_state.partition_bookdeg=_out_degrees(local_g,partition_book,self.n)returnOutDegreeResponse(deg,self.order_id)classOutDegreeResponse(Response):"""The response for out-degree"""def__init__(self,deg,order_id):self.val=degself.order_id=order_iddef__setstate__(self,state):self.val,self.order_id=statedef__getstate__(self):returnself.val,self.order_idclassInSubgraphRequest(Request):"""InSubgraph Request"""def__init__(self,nodes):self.seed_nodes=nodesdef__setstate__(self,state):self.seed_nodes=statedef__getstate__(self):returnself.seed_nodesdefprocess_request(self,server_state):local_g=server_state.graphpartition_book=server_state.partition_bookglobal_src,global_dst,global_eids=_in_subgraph(local_g,partition_book,self.seed_nodes)returnSubgraphResponse(global_src,global_dst,global_eids=global_eids)defmerge_graphs(res_list,num_nodes,exclude_edges=None):"""Merge request from multiple servers"""iflen(res_list)>1:srcs=[]dsts=[]eids=[]etype_ids=[]forresinres_list:srcs.append(res.global_src)dsts.append(res.global_dst)eids.append(res.global_eids)etype_ids.append(res.etype_ids)src_tensor=F.cat(srcs,0)dst_tensor=F.cat(dsts,0)eid_tensor=Noneifeids[0]isNoneelseF.cat(eids,0)etype_id_tensor=Noneifetype_ids[0]isNoneelseF.cat(etype_ids,0)else:src_tensor=res_list[0].global_srcdst_tensor=res_list[0].global_dsteid_tensor=res_list[0].global_eidsetype_id_tensor=res_list[0].etype_idsifexclude_edgesisnotNone:mask=torch.isin(eid_tensor,exclude_edges,assume_unique=True,invert=True)src_tensor=src_tensor[mask]dst_tensor=dst_tensor[mask]eid_tensor=eid_tensor[mask]ifetype_id_tensorisnotNone:etype_id_tensor=etype_id_tensor[mask]g=graph((src_tensor,dst_tensor),num_nodes=num_nodes)ifeid_tensorisnotNone:g.edata[EID]=eid_tensorifetype_id_tensorisnotNone:g.edata[ETYPE]=etype_id_tensorreturngLocalSampledGraph=namedtuple(# pylint: disable=unexpected-keyword-arg"LocalSampledGraph","global_src global_dst global_eids etype_ids",defaults=(None,None,None,None),)def_distributed_access(g,nodes,issue_remote_req,local_access,exclude_edges=None):"""A routine that fetches local neighborhood of nodes from the distributed graph. The local neighborhood of some nodes are stored in the local machine and the other nodes have their neighborhood on remote machines. This code will issue remote access requests first before fetching data from the local machine. In the end, we combine the data from the local machine and remote machines. In this way, we can hide the latency of accessing data on remote machines. Parameters ---------- g : DistGraph The distributed graph nodes : tensor The nodes whose neighborhood are to be fetched. issue_remote_req : callable The function that issues requests to access remote data. local_access : callable The function that reads data on the local machine. exclude_edges : tensor The edges to exclude after sampling. Returns ------- DGLGraph The subgraph that contains the neighborhoods of all input nodes. """req_list=[]partition_book=g.get_partition_book()ifnotisinstance(nodes,torch.Tensor):nodes=toindex(nodes).tousertensor()partition_id=partition_book.nid2partid(nodes)local_nids=Noneforpidinrange(partition_book.num_partitions()):node_id=F.boolean_mask(nodes,partition_id==pid)# We optimize the sampling on a local partition if the server and the client# run on the same machine. With a good partitioning, most of the seed nodes# should reside in the local partition. If the server and the client# are not co-located, the client doesn't have a local partition.ifpid==partition_book.partidandg.local_partitionisnotNone:assertlocal_nidsisNonelocal_nids=node_ideliflen(node_id)!=0:req=issue_remote_req(node_id)req_list.append((pid,req))# send requests to the remote machine.msgseq2pos=Noneiflen(req_list)>0:msgseq2pos=send_requests_to_machine(req_list)# sample neighbors for the nodes in the local partition.res_list=[]iflocal_nidsisnotNone:res=local_access(g.local_partition,partition_book,local_nids)res_list.append(res)# receive responses from remote machines.ifmsgseq2posisnotNone:results=recv_responses(msgseq2pos)res_list.extend(results)sampled_graph=merge_graphs(res_list,g.num_nodes(),exclude_edges=exclude_edges)returnsampled_graphdef_frontier_to_heterogeneous_graph(g,frontier,gpb):# We need to handle empty frontiers correctly.iffrontier.num_edges()==0:data_dict={etype:(np.zeros(0),np.zeros(0))foretypeing.canonical_etypes}returnheterograph(data_dict,{ntype:g.num_nodes(ntype)forntypeing.ntypes},idtype=g.idtype,)# For DGL partitions, the global edge IDs are always stored in the edata.# For GraphBolt partitions, the edge type IDs are always stored in the# edata. As for the edge IDs, they are stored in the edata if the graph is# partitioned with `store_eids=True`. Otherwise, the edge IDs are not# stored.etype_ids,type_wise_eids=(gpb.map_to_per_etype(frontier.edata[EID])ifEIDinfrontier.edataelse(frontier.edata[ETYPE],None))etype_ids,idx=F.sort_1d(etype_ids)iftype_wise_eidsisnotNone:type_wise_eids=F.gather_row(type_wise_eids,idx)# Sort the edges by their edge types.src,dst=frontier.edges()src,dst=F.gather_row(src,idx),F.gather_row(dst,idx)src_ntype_ids,src=gpb.map_to_per_ntype(src)dst_ntype_ids,dst=gpb.map_to_per_ntype(dst)data_dict=dict()edge_ids={}foretid,etypeinenumerate(g.canonical_etypes):src_ntype,_,dst_ntype=etypesrc_ntype_id=g.get_ntype_id(src_ntype)dst_ntype_id=g.get_ntype_id(dst_ntype)type_idx=etype_ids==etiddata_dict[etype]=(F.boolean_mask(src,type_idx),F.boolean_mask(dst,type_idx),)if"DGL_DIST_DEBUG"inos.environ:asserttorch.all(src_ntype_id==src_ntype_ids[type_idx]),"source ntype is is not expected."asserttorch.all(dst_ntype_id==dst_ntype_ids[type_idx]),"destination ntype is is not expected."iftype_wise_eidsisnotNone:edge_ids[etype]=F.boolean_mask(type_wise_eids,type_idx)hg=heterograph(data_dict,{ntype:g.num_nodes(ntype)forntypeing.ntypes},idtype=g.idtype,)foretypeinedge_ids:hg.edges[etype].data[EID]=edge_ids[etype]returnhg
[docs]defsample_etype_neighbors(g,nodes,fanout,edge_dir="in",prob=None,exclude_edges=None,replace=False,etype_sorted=True,use_graphbolt=False,):"""Sample from the neighbors of the given nodes from a distributed graph. For each node, a number of inbound (or outbound when ``edge_dir == 'out'``) edges will be randomly chosen. The returned graph will contain all the nodes in the original graph, but only the sampled edges. Node/edge features are not preserved. The original IDs of the sampled edges are stored as the `dgl.EID` feature in the returned graph. This function assumes the input is a homogeneous ``DGLGraph`` with the edges ordered by their edge types. The sampled subgraph is also stored in the homogeneous graph format. That is, all nodes and edges are assigned with unique IDs (in contrast, we typically use a type name and a node/edge ID to identify a node or an edge in ``DGLGraph``). We refer to this type of IDs as *homogeneous ID*. Users can use :func:`dgl.distributed.GraphPartitionBook.map_to_per_ntype` and :func:`dgl.distributed.GraphPartitionBook.map_to_per_etype` to identify their node/edge types and node/edge IDs of that type. Parameters ---------- g : DistGraph The distributed graph.. nodes : tensor or dict Node IDs to sample neighbors from. If it's a dict, it should contain only one key-value pair to make this API consistent with dgl.sampling.sample_neighbors. fanout : int or dict[etype, int] The number of edges to be sampled for each node per edge type. If an integer is given, DGL assumes that the same fanout is applied to every edge type. If -1 is given, all of the neighbors will be selected. edge_dir : str, optional Determines whether to sample inbound or outbound edges. Can take either ``in`` for inbound edges or ``out`` for outbound edges. prob : str, optional Feature name used as the (unnormalized) probabilities associated with each neighboring edge of a node. The feature must have only one element for each edge. The features must be non-negative floats, and the sum of the features of inbound/outbound edges for every node must be positive (though they don't have to sum up to one). Otherwise, the result will be undefined. exclude_edges : tensor, optional The edges to exclude when sampling. Homogeneous edge IDs are used. replace : bool, optional If True, sample with replacement. When sampling with replacement, the sampled subgraph could have parallel edges. For sampling without replacement, if fanout > the number of neighbors, all the neighbors are sampled. If fanout == -1, all neighbors are collected. etype_sorted : bool, optional Indicates whether etypes are sorted. use_graphbolt : bool, optional Whether to use GraphBolt for sampling. Returns ------- DGLGraph A sampled subgraph containing only the sampled neighboring edges. It is on CPU. """ifisinstance(fanout,int):fanout=F.full_1d(len(g.canonical_etypes),fanout,F.int64,F.cpu())else:etype_ids={etype:ifori,etypeinenumerate(g.canonical_etypes)}fanout_array=[None]*len(g.canonical_etypes)foretype,vinfanout.items():c_etype=g.to_canonical_etype(etype)fanout_array[etype_ids[c_etype]]=vassertall(visnotNoneforvinfanout_array),("Not all etypes have valid fanout. Please make sure passed-in ""fanout in dict includes all the etypes in graph. Passed-in "f"fanout: {fanout}, graph etypes: {g.canonical_etypes}.")fanout=F.tensor(fanout_array,dtype=F.int64)gpb=g.get_partition_book()ifisinstance(nodes,dict):homo_nids=[]forntypeinnodes.keys():assert(ntypeing.ntypes),"The sampled node type {} does not exist in the input graph".format(ntype)ifF.is_tensor(nodes[ntype]):typed_nodes=nodes[ntype]else:typed_nodes=toindex(nodes[ntype]).tousertensor()homo_nids.append(gpb.map_to_homo_nid(typed_nodes,ntype))nodes=F.cat(homo_nids,0)defissue_remote_req(node_ids):ifprobisnotNoneand(notuse_graphbolt):# See NOTE 1_prob=[(# NOTE (BarclayII)# Currently DistGraph.edges[] does not accept canonical etype.g.edges[etype].data[prob].kvstore_keyifprobing.edges[etype].dataelse"")foretypeing.canonical_etypes]else:_prob=probreturnSamplingRequestEtype(node_ids,fanout,edge_dir=edge_dir,prob=_prob,exclude_edges=None,replace=replace,etype_sorted=etype_sorted,use_graphbolt=use_graphbolt,)deflocal_access(local_g,partition_book,local_nids):etype_offset=gpb.local_etype_offset# See NOTE 1ifprobisnotNoneand(notuse_graphbolt):_prob=[(g.edges[etype].data[prob].local_partitionifprobing.edges[etype].dataelseNone)foretypeing.canonical_etypes]else:_prob=probreturn_sample_etype_neighbors(use_graphbolt,local_g,partition_book,local_nids,fanout,edge_dir=edge_dir,prob=_prob,exclude_edges=None,replace=replace,etype_offset=etype_offset,etype_sorted=etype_sorted,)frontier=_distributed_access(g,nodes,issue_remote_req,local_access,exclude_edges=exclude_edges)ifnotgpb.is_homogeneous:return_frontier_to_heterogeneous_graph(g,frontier,gpb)else:returnfrontier
[docs]defsample_neighbors(g,nodes,fanout,edge_dir="in",prob=None,exclude_edges=None,replace=False,use_graphbolt=False,):"""Sample from the neighbors of the given nodes from a distributed graph. For each node, a number of inbound (or outbound when ``edge_dir == 'out'``) edges will be randomly chosen. The returned graph will contain all the nodes in the original graph, but only the sampled edges. Node/edge features are not preserved. The original IDs of the sampled edges are stored as the `dgl.EID` feature in the returned graph. For heterogeneous graphs, ``nodes`` is a dictionary whose key is node type and the value is type-specific node IDs. Parameters ---------- g : DistGraph The distributed graph.. nodes : tensor or dict Node IDs to sample neighbors from. If it's a dict, it should contain only one key-value pair to make this API consistent with dgl.sampling.sample_neighbors. fanout : int The number of edges to be sampled for each node. If -1 is given, all of the neighbors will be selected. edge_dir : str, optional Determines whether to sample inbound or outbound edges. Can take either ``in`` for inbound edges or ``out`` for outbound edges. prob : str, optional Feature name used as the (unnormalized) probabilities associated with each neighboring edge of a node. The feature must have only one element for each edge. The features must be non-negative floats, and the sum of the features of inbound/outbound edges for every node must be positive (though they don't have to sum up to one). Otherwise, the result will be undefined. exclude_edges: tensor or dict, optional Edge IDs to exclude during sampling neighbors for the seed nodes. This argument can take a single ID tensor or a dictionary of edge types and ID tensors. If a single tensor is given, the graph must only have one type of nodes. replace : bool, optional If True, sample with replacement. When sampling with replacement, the sampled subgraph could have parallel edges. For sampling without replacement, if fanout > the number of neighbors, all the neighbors are sampled. If fanout == -1, all neighbors are collected. use_graphbolt : bool, optional Whether to use GraphBolt for sampling. Returns ------- DGLGraph A sampled subgraph containing only the sampled neighboring edges. It is on CPU. """gpb=g.get_partition_book()ifnotgpb.is_homogeneous:assertisinstance(nodes,dict)homo_nids=[]forntypeinnodes:assert(ntypeing.ntypes),"The sampled node type does not exist in the input graph"ifF.is_tensor(nodes[ntype]):typed_nodes=nodes[ntype]else:typed_nodes=toindex(nodes[ntype]).tousertensor()homo_nids.append(gpb.map_to_homo_nid(typed_nodes,ntype))nodes=F.cat(homo_nids,0)elifisinstance(nodes,dict):assertlen(nodes)==1nodes=list(nodes.values())[0]defissue_remote_req(node_ids):ifprobisnotNoneand(notuse_graphbolt):# See NOTE 1_prob=g.edata[prob].kvstore_keyelse:_prob=probreturnSamplingRequest(node_ids,fanout,edge_dir=edge_dir,prob=_prob,exclude_edges=None,replace=replace,use_graphbolt=use_graphbolt,)deflocal_access(local_g,partition_book,local_nids):# See NOTE 1_prob=([g.edata[prob].local_partition]ifprobisnotNoneand(notuse_graphbolt)elseprob)return_sample_neighbors(use_graphbolt,local_g,partition_book,local_nids,fanout,edge_dir=edge_dir,prob=_prob,exclude_edges=None,replace=replace,)frontier=_distributed_access(g,nodes,issue_remote_req,local_access,exclude_edges=exclude_edges)ifnotgpb.is_homogeneous:return_frontier_to_heterogeneous_graph(g,frontier,gpb)else:returnfrontier
def_distributed_edge_access(g,edges,issue_remote_req,local_access):"""A routine that fetches local edges from distributed graph. The source and destination nodes of local edges are stored in the local machine and others are stored on remote machines. This code will issue remote access requests first before fetching data from the local machine. In the end, we combine the data from the local machine and remote machines. Parameters ---------- g : DistGraph The distributed graph edges : tensor The edges to find their source and destination nodes. issue_remote_req : callable The function that issues requests to access remote data. local_access : callable The function that reads data on the local machine. Returns ------- tensor The source node ID array. tensor The destination node ID array. """req_list=[]partition_book=g.get_partition_book()edges=toindex(edges).tousertensor()partition_id=partition_book.eid2partid(edges)local_eids=Nonereorder_idx=[]forpidinrange(partition_book.num_partitions()):mask=partition_id==pidedge_id=F.boolean_mask(edges,mask)reorder_idx.append(F.nonzero_1d(mask))ifpid==partition_book.partidandg.local_partitionisnotNone:assertlocal_eidsisNonelocal_eids=edge_ideliflen(edge_id)!=0:req=issue_remote_req(edge_id,pid)req_list.append((pid,req))# send requests to the remote machine.msgseq2pos=Noneiflen(req_list)>0:msgseq2pos=send_requests_to_machine(req_list)# handle edges in local partition.src_ids=F.zeros_like(edges)dst_ids=F.zeros_like(edges)iflocal_eidsisnotNone:src,dst=local_access(g.local_partition,partition_book,local_eids)src_ids=F.scatter_row(src_ids,reorder_idx[partition_book.partid],src)dst_ids=F.scatter_row(dst_ids,reorder_idx[partition_book.partid],dst)# receive responses from remote machines.ifmsgseq2posisnotNone:results=recv_responses(msgseq2pos)forresultinresults:src=result.global_srcdst=result.global_dstsrc_ids=F.scatter_row(src_ids,reorder_idx[result.order_id],src)dst_ids=F.scatter_row(dst_ids,reorder_idx[result.order_id],dst)returnsrc_ids,dst_ids
[docs]deffind_edges(g,edge_ids):"""Given an edge ID array, return the source and destination node ID array ``s`` and ``d`` from a distributed graph. ``s[i]`` and ``d[i]`` are source and destination node ID for edge ``eid[i]``. Parameters ---------- g : DistGraph The distributed graph. edges : tensor The edge ID array. Returns ------- tensor The source node ID array. tensor The destination node ID array. """defissue_remote_req(edge_ids,order_id):returnEdgesRequest(edge_ids,order_id)deflocal_access(local_g,partition_book,edge_ids):return_find_edges(local_g,partition_book,edge_ids)return_distributed_edge_access(g,edge_ids,issue_remote_req,local_access)
[docs]defin_subgraph(g,nodes):"""Return the subgraph induced on the inbound edges of the given nodes. The subgraph keeps the same type schema and all the nodes are preserved regardless of whether they have an edge or not. Node/edge features are not preserved. The original IDs of the extracted edges are stored as the `dgl.EID` feature in the returned graph. For now, we only support the input graph with one node type and one edge type. Parameters ---------- g : DistGraph The distributed graph structure. nodes : tensor or dict Node ids to sample neighbors from. Returns ------- DGLGraph The subgraph. One can retrieve the mapping from subgraph edge ID to parent edge ID via ``dgl.EID`` edge features of the subgraph. """ifisinstance(nodes,dict):assert(len(nodes)==1),"The distributed in_subgraph only supports one node type for now."nodes=list(nodes.values())[0]defissue_remote_req(node_ids):returnInSubgraphRequest(node_ids)deflocal_access(local_g,partition_book,local_nids):return_in_subgraph(local_g,partition_book,local_nids)return_distributed_access(g,nodes,issue_remote_req,local_access)
def_distributed_get_node_property(g,n,issue_remote_req,local_access):req_list=[]partition_book=g.get_partition_book()n=toindex(n).tousertensor()partition_id=partition_book.nid2partid(n)local_nids=Nonereorder_idx=[]forpidinrange(partition_book.num_partitions()):mask=partition_id==pidnid=F.boolean_mask(n,mask)reorder_idx.append(F.nonzero_1d(mask))ifpid==partition_book.partidandg.local_partitionisnotNone:assertlocal_nidsisNonelocal_nids=nideliflen(nid)!=0:req=issue_remote_req(nid,pid)req_list.append((pid,req))# send requests to the remote machine.msgseq2pos=Noneiflen(req_list)>0:msgseq2pos=send_requests_to_machine(req_list)# handle edges in local partition.vals=Noneiflocal_nidsisnotNone:local_vals=local_access(g.local_partition,partition_book,local_nids)shape=list(F.shape(local_vals))shape[0]=len(n)vals=F.zeros(shape,F.dtype(local_vals),F.cpu())vals=F.scatter_row(vals,reorder_idx[partition_book.partid],local_vals)# receive responses from remote machines.ifmsgseq2posisnotNone:results=recv_responses(msgseq2pos)iflen(results)>0andvalsisNone:shape=list(F.shape(results[0].val))shape[0]=len(n)vals=F.zeros(shape,F.dtype(results[0].val),F.cpu())forresultinresults:val=result.valvals=F.scatter_row(vals,reorder_idx[result.order_id],val)returnvalsdefin_degrees(g,v):"""Get in-degrees"""defissue_remote_req(v,order_id):returnInDegreeRequest(v,order_id)deflocal_access(local_g,partition_book,v):return_in_degrees(local_g,partition_book,v)return_distributed_get_node_property(g,v,issue_remote_req,local_access)defout_degrees(g,u):"""Get out-degrees"""defissue_remote_req(u,order_id):returnOutDegreeRequest(u,order_id)deflocal_access(local_g,partition_book,u):return_out_degrees(local_g,partition_book,u)return_distributed_get_node_property(g,u,issue_remote_req,local_access)register_service(SAMPLING_SERVICE_ID,SamplingRequest,SubgraphResponse)register_service(EDGES_SERVICE_ID,EdgesRequest,FindEdgeResponse)register_service(INSUBGRAPH_SERVICE_ID,InSubgraphRequest,SubgraphResponse)register_service(OUTDEGREE_SERVICE_ID,OutDegreeRequest,OutDegreeResponse)register_service(INDEGREE_SERVICE_ID,InDegreeRequest,InDegreeResponse)register_service(ETYPE_SAMPLING_SERVICE_ID,SamplingRequestEtype,SubgraphResponse)