Repository: dmlc/dgl
Pull request: [DIST] Message size to retrieve SHUFFLE_GLOBAL_NIDs is resulting in very large messages and resulting in killed process
## Description
<!-- Brief description. Refer to the related issues if existed.
It'll be great if relevant reviewers can be assigned as well.-->
When retrieving SHUFFLE_GLOBAL_NID's for src and dst end points of the local edges in the current graph partition we used to send both the source and destination end points as a single list to the distributed lookup service. For Fe_large graph this message size is becoming about 50GB+ (3.2 billion edges). This increases the pressure on the memory used to send/receive such a large message and in the case of fe_large graph the process is getting killed. 

Now, we batch the outgoing message to distributed lookup service and ensure that the in each iteration we only send a message of maximum size of 250 MILLION entries (which is about 2 GB). This will reduce the memory need to support this operation and with this change fe_large graph is able to run correctly. 

Also another optimization is to NOT send destination end points to the distributed lookup service, since current process will have associated SHUFFLE_GLOBAL_NID's. We use current process local information to retrieve these IDs. 

## Checklist
Please feel free to remove inapplicable items for your PR.
- [X] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
- [X] I've leverage the [tools](https://docs.google.com/document/d/1iHyj7zlmygKSk5gBPsqIqL5ASPzJSPREaNT_QdsiYA4/edit) to beautify the python and c++ code.
- [X] The PR is complete and small, read the [Google eng practice (CL equals to PR)](https://google.github.io/eng-practices/review/developer/small-cls.html) to understand more about small PR. In DGL, we consider PRs with less than 200 lines of core code change are small (example, test and documentation could be exempted).
- [X] All changes have test coverage
- [X] Code is well-documented
- [X] To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
- [ ] Related issue is referred in this PR
- [ ] If the PR is for a new model/paper, I've updated the example index [here](../examples/README.md).

## Changes
<!-- You could use following template
- [ ] Feature1, tests, (and when applicable, API doc)
- [ ] Feature2, tests, (and when applicable, API doc)
-->


Problem statement:
Memory optimizations for Fe-Large graph when reading large files using pyarrow
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

Reading large csv files using pyarrow is resulting in very large memory footprint for the python pipeline process. Even after the reading is complete this memory is lingering around and not getting released back to the system. 

## To Reproduce

Steps to reproduce the behavior:

1. Run pipeline with Fe_Large graph and notice the memory usage using either vmstat/top or any other memory related system commands. 

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

 - DGL Version (e.g., 1.0):
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3):
 - OS (e.g., Linux):
 - How you installed DGL (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->

Hints from issue discussion:

Commit message: Update function signature for allgather_sizes function call.

Patch to review:
diff --git a/tools/distpartitioning/globalids.py b/tools/distpartitioning/globalids.py
index ca7236c6d869..c3698b9e7559 100644
--- a/tools/distpartitioning/globalids.py
+++ b/tools/distpartitioning/globalids.py
@@ -86,13 +86,47 @@ def lookup_shuffle_global_nids_edges(rank, world_size, edge_data, id_lookup, nod
         dictionary where keys are column names and values are numpy arrays representing all the
         edges present in the current graph partition
     '''
+    # Make sure that the outgoing message size does not exceed 2GB in size. 
+    # Even though gloo can handle upto 10GB size of data in the outgoing messages,
+    # it needs additional memory to store temporary information into the buffers which will increase
+    # the memory needs of the process. 
+    MILLION = 1000 * 1000
+    BATCH_SIZE = 250 * MILLION
     memory_snapshot("GlobalToShuffleIDMapBegin: ", rank)
-    node_list = np.concatenate([edge_data[constants.GLOBAL_SRC_ID], edge_data[constants.GLOBAL_DST_ID]])
-    shuffle_ids = id_lookup.get_shuffle_nids(node_list,
+    node_list = edge_data[constants.GLOBAL_SRC_ID]
+
+    # Determine the no. of times each process has to send alltoall messages.
+    all_sizes = allgather_sizes([node_list.shape[0]], world_size, return_sizes=True)
+    max_count = np.amax(all_sizes)
+    if max_count >= BATCH_SIZE:
+        num_splits = max_count // BATCH_SIZE + 1
+    else:
+        num_splits = 1
+
+    # Split the message into batches and send.
+    splits = np.array_split(node_list, num_splits)
+    shuffle_mappings = []
+    for item in splits:
+        logging.info(f'[Rank: {rank} SENDING : {item.shape}')
+        shuffle_ids = id_lookup.get_shuffle_nids(item,
                                             node_data[constants.GLOBAL_NID],
                                             node_data[constants.SHUFFLE_GLOBAL_NID])
+        shuffle_mappings.append(shuffle_ids)
+
+    shuffle_ids = np.concatenate(shuffle_mappings)
+    assert shuffle_ids.shape[0] == node_list.shape[0]
+    edge_data[constants.SHUFFLE_GLOBAL_SRC_ID] = shuffle_ids
+
+    # Destination end points of edges are owned by the current node and therefore
+    # should have corresponding SHUFFLE_GLOBAL_NODE_IDs. 
+    # Here retrieve SHUFFLE_GLOBAL_NODE_IDs for the destination end points of local edges.
+    uniq_ids, inverse_idx = np.unique(edge_data[constants.GLOBAL_DST_ID], return_inverse=True)
+    common, idx1, idx2 = np.intersect1d(uniq_ids, node_data[constants.GLOBAL_NID], assume_unique=True, return_indices=True)
+    assert len(common) == len(uniq_ids)
+
+    edge_data[constants.SHUFFLE_GLOBAL_DST_ID] = node_data[constants.SHUFFLE_GLOBAL_NID][idx2][inverse_idx]
+    assert len(edge_data[constants.SHUFFLE_GLOBAL_DST_ID]) == len(edge_data[constants.GLOBAL_DST_ID])
 
-    edge_data[constants.SHUFFLE_GLOBAL_SRC_ID], edge_data[constants.SHUFFLE_GLOBAL_DST_ID] = np.split(shuffle_ids, 2)
     memory_snapshot("GlobalToShuffleIDMap_AfterLookupServiceCalls: ", rank)
     return edge_data
 
diff --git a/tools/distpartitioning/gloo_wrapper.py b/tools/distpartitioning/gloo_wrapper.py
index f6dbf1449209..7af21527d00c 100644
--- a/tools/distpartitioning/gloo_wrapper.py
+++ b/tools/distpartitioning/gloo_wrapper.py
@@ -2,7 +2,7 @@
 import torch
 import torch.distributed as dist
 
-def allgather_sizes(send_data, world_size):
+def allgather_sizes(send_data, world_size, return_sizes=False):
     """ 
     Perform all gather on list lengths, used to compute prefix sums
     to determine the offsets on each ranks. This is used to allocate
@@ -14,6 +14,9 @@ def allgather_sizes(send_data, world_size):
         Data on which allgather is performed.
     world_size : integer
         No. of processes configured for execution
+    return_sizes : bool
+        Boolean flag to indicate whether to return raw sizes from each process
+        or perform prefix sum on the raw sizes.
 
     Returns : 
     ---------
@@ -30,6 +33,10 @@ def allgather_sizes(send_data, world_size):
     #all_gather message
     dist.all_gather(in_tensor, out_tensor)
 
+    # Return on the raw sizes from each process
+    if return_sizes:
+        return torch.cat(in_tensor).numpy()
+
     #gather sizes in on array to return to the invoking function
     rank_sizes = np.zeros(world_size + 1, dtype=np.int64)
     count = rank_sizes[0]


Code review comments:
diff_hunk: @@ -86,13 +86,47 @@ def lookup_shuffle_global_nids_edges(rank, world_size, edge_data, id_lookup, nod
         dictionary where keys are column names and values are numpy arrays representing all the
         edges present in the current graph partition
     '''
+    # Make sure that the outgoing message size does not exceed 2GB in size. 
+    # Even though gloo can handle upto 10GB size of data in the outgoing messages,
+    # it needs additional memory to store temporary information into the buffers which will increase
+    # the memory needs of the process. 
+    MILLION = 1000 * 1000
+    BATCH_SIZE = 250 * MILLION
     memory_snapshot("GlobalToShuffleIDMapBegin: ", rank)
-    node_list = np.concatenate([edge_data[constants.GLOBAL_SRC_ID], edge_data[constants.GLOBAL_DST_ID]])
-    shuffle_ids = id_lookup.get_shuffle_nids(node_list,
+    node_list = edge_data[constants.GLOBAL_SRC_ID]
+
+    # Determine the no. of times each process has to send alltoall messages.
+    all_sizes = allgather_sizes([node_list.shape[0]], world_size, return_sizes=True)
+    max_count = np.amax(all_sizes)
+    if max_count >= BATCH_SIZE:
path: tools/distpartitioning/globalids.py
text: @user1:
`num_splits = max_count // BATCH_SIZE + 1` should be ok regardless of `max_count % BATCH_SIZE == 0` or `max_count >= BATCH_SIZE`.

@author:
Done. Updated the code accordingly.

diff_hunk: @@ -86,13 +86,47 @@ def lookup_shuffle_global_nids_edges(rank, world_size, edge_data, id_lookup, nod
         dictionary where keys are column names and values are numpy arrays representing all the
         edges present in the current graph partition
     '''
+    # Make sure that the outgoing message size does not exceed 2GB in size. 
+    # Even though gloo can handle upto 10GB size of data in the outgoing messages,
+    # it needs additional memory to store temporary information into the buffers which will increase
+    # the memory needs of the process. 
+    MILLION = 1000 * 1000
+    BATCH_SIZE = 250 * MILLION
     memory_snapshot("GlobalToShuffleIDMapBegin: ", rank)
-    node_list = np.concatenate([edge_data[constants.GLOBAL_SRC_ID], edge_data[constants.GLOBAL_DST_ID]])
-    shuffle_ids = id_lookup.get_shuffle_nids(node_list,
+    node_list = edge_data[constants.GLOBAL_SRC_ID]
+
+    # Determine the no. of times each process has to send alltoall messages.
+    all_sizes = allgather_sizes([node_list.shape[0]], world_size, return_sizes=True)
+    max_count = np.amax(all_sizes)
+    if max_count >= BATCH_SIZE:
+        num_splits = max_count // BATCH_SIZE + 1
+    else:
+        num_splits = 1
+
+    # Split the message into batches and send.
+    splits = np.array_split(node_list, num_splits)
+    shuffle_mappings = []
+    for item in splits:
+        logging.info(f'[Rank: {rank} SENDING : {item.shape}')
path: tools/distpartitioning/globalids.py
text: @user1:
`NameError: name 'logging' is not defined`

@author:
Removed the logging.info line. This was only for debugging purposes.