Knowing Halo nodes in a partitioned graph

[emduc]

Hi,

I partitioned a HeteroData graph using distributed.Partitioner and generate_partition(). I noticed that in each partition, I only have the local node, that is without any halo nodes or data replication. I thought that there would be duplication of the halo nodes in each partition, but it does not seem to be the case. How can I find a way to retrieve all my halo nodes for a given partition?

As a reminder, this is how the partitioning looks like:

root/
|-- META.json
|-- node_map.pt
|-- edge_map.pt
|-- part0/
|-- graph.pt
|-- node_feats.pt
|-- edge_feats.pt
|-- part1/
|-- graph.pt
|-- node_feats.pt
|-- edge_feats.pt

My goal here is to only do message passing in regard to the nodes shared between two partitions (halo nodes).
Is there any straight forward way to get my halo nodes?

[kgajdamo]

Hi @emduc!
The graph.pt file contains information about all node connections in a given partition and therefore also the halo nodes connections. graph.pt has the following fields:

'edge_id': edge_id,
'row': global_row,
'col': global_col,
'size': (data.num_nodes, data.num_nodes),

What will be interesting for you are row and col, which are the indices of the src and dst nodes.
There are no files that directly highlight the halo nodes, so you have to obtain them yourself based on the values from row/col and from node_map.pt, which contains information about the assignment of all nodes to individual partitions.
An example of obtaining information about halo nodes in a situation when we have two partitions could be as follows:

node_pb = torch.load(osp.join(root_dir, 'node_map.pt'))
graph_data = torch.load(osp.join(part_dir, 'graph.pt'))

# Assume machine 0 is local
p1_id = 1
# Get all node indices of partition 1:
p1_ids = (node_pb == p1_id).nonzero().squeeze().numpy()
row = graph_data['row'].numpy()
col = graph_data['col'].numpy()

# Find indices of nodes from partition 1 in "row":
indices = np.where(np.in1d(row, p1_ids))[0]
# Should be empty (here are only local node ids):
assert(np.where(np.in1d(col, p1_ids))[0].size == 0)
# Take the corresponding values from the indices from "row" and "col":
halo_nodes_0 = col[indices]
halo_nodes_1 = row[indices]

# Should be empty:
assert(np.where(np.in1d(halo_nodes_0, halo_nodes_1))[0].size == 0)

# Partition 0 halo nodes:
halo_nodes_0 = torch.from_numpy(halo_nodes_0)
# Partition 1 halo nodes:
halo_nodes_1 = torch.from_numpy(halo_nodes_1)

If any more questions arises feel free to ask :)