Repository: pyg-team/pytorch_geometric
Pull request #7974 (2023-09-04 13:58:34)

Add `DistNeighborSampler`
This code belongs to the part of the whole distributed training for PyG.

`DistNeighborSampler` leverages the `NeighborSampler` class from `pytorch_geometric` and the `neighbor_sample` function from `pyg-lib`. However, due to the fact that in case of distributed training it is required to synchronise the results between machines after each layer, the part of the code responsible for sampling was implemented in python.

Added suport for the following sampling methods:
- node, edge, negative, disjoint, temporal

**TODOs:**

- [x] finish hetero part
- [x] subgraph sampling

**This PR should be merged together with other distributed PRs:**
pyg-lib: [#246](https://github.com/pyg-team/pyg-lib/pull/246), [#252](https://github.com/pyg-team/pyg-lib/pull/252)
GraphStore\FeatureStore: https://github.com/pyg-team/pytorch_geometric/pull/8083
DistLoaders:
1.  https://github.com/pyg-team/pytorch_geometric/pull/8079
2.  https://github.com/pyg-team/pytorch_geometric/pull/8080
3.  https://github.com/pyg-team/pytorch_geometric/pull/8085

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7974](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (66623fc) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/cf60b5fdf2ef38dd594946acbf392ed7678d1226?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (cf60b5f) will **decrease** coverage by `1.34%`.
> Report is 3 commits behind head on master.
> The diff coverage is `14.04%`.

> :exclamation: Current head 66623fc differs from pull request most recent head 6fe0f67. Consider uploading reports for the commit 6fe0f67 to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #7974      +/-   ##
==========================================
- Coverage   89.58%   88.25%   -1.34%     
==========================================
  Files         459      463       +4     
  Lines       26951    27400     +449     
==========================================
+ Hits        24145    24181      +36     
- Misses       2806     3219     +413     
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [...rch\_geometric/distributed/dist\_neighbor\_sampler.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL2Rpc3RfbmVpZ2hib3Jfc2FtcGxlci5weQ==) | `0.00% <0.00%> (ø)` | |
| [torch\_geometric/distributed/event\_loop.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL2V2ZW50X2xvb3AucHk=) | `0.00% <0.00%> (ø)` | |
| [torch\_geometric/distributed/local\_graph\_store.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL2xvY2FsX2dyYXBoX3N0b3JlLnB5) | `91.20% <ø> (ø)` | |
| [torch\_geometric/distributed/utils.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL3V0aWxzLnB5) | `0.00% <0.00%> (ø)` | |
| [torch\_geometric/sampler/base.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3NhbXBsZXIvYmFzZS5weQ==) | `89.76% <27.27%> (-5.93%)` | :arrow_down: |
| [torch\_geometric/sampler/neighbor\_sampler.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3NhbXBsZXIvbmVpZ2hib3Jfc2FtcGxlci5weQ==) | `86.47% <61.22%> (-6.33%)` | :arrow_down: |
| [torch\_geometric/data/graph\_store.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZ3JhcGhfc3RvcmUucHk=) | `93.98% <78.78%> (-4.38%)` | :arrow_down: |
| [torch\_geometric/sampler/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3NhbXBsZXIvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |

... and [5 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7974/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4b4f7e947..f0be4b972 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -75,7 +75,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added a CPU-based and GPU-based `map_index` implementation ([#7493](https://github.com/pyg-team/pytorch_geometric/pull/7493), [#7764](https://github.com/pyg-team/pytorch_geometric/pull/7764) [#7765](https://github.com/pyg-team/pytorch_geometric/pull/7765))
 - Added the `AmazonBook` heterogeneous dataset ([#7483](https://github.com/pyg-team/pytorch_geometric/pull/7483))
 - Added hierarchical heterogeneous GraphSAGE example on OGB-MAG ([#7425](https://github.com/pyg-team/pytorch_geometric/pull/7425))
-- Added the `torch_geometric.distributed` package ([#7451](https://github.com/pyg-team/pytorch_geometric/pull/7451), [#7452](https://github.com/pyg-team/pytorch_geometric/pull/7452)), [#7482](https://github.com/pyg-team/pytorch_geometric/pull/7482), [#7502](https://github.com/pyg-team/pytorch_geometric/pull/7502), [#7628](https://github.com/pyg-team/pytorch_geometric/pull/7628), [#7671](https://github.com/pyg-team/pytorch_geometric/pull/7671), [#7846](https://github.com/pyg-team/pytorch_geometric/pull/7846), [#7715](https://github.com/pyg-team/pytorch_geometric/pull/7715))
+- Added the `torch_geometric.distributed` package ([#7451](https://github.com/pyg-team/pytorch_geometric/pull/7451), [#7452](https://github.com/pyg-team/pytorch_geometric/pull/7452)), [#7482](https://github.com/pyg-team/pytorch_geometric/pull/7482), [#7502](https://github.com/pyg-team/pytorch_geometric/pull/7502), [#7628](https://github.com/pyg-team/pytorch_geometric/pull/7628), [#7671](https://github.com/pyg-team/pytorch_geometric/pull/7671), [#7846](https://github.com/pyg-team/pytorch_geometric/pull/7846), [#7715](https://github.com/pyg-team/pytorch_geometric/pull/7715), [#7974](https://github.com/pyg-team/pytorch_geometric/pull/7974))
 - Added the `GDELTLite` dataset ([#7442](https://github.com/pyg-team/pytorch_geometric/pull/7442))
 - Added the `approx_knn` function for approximated nearest neighbor search ([#7421](https://github.com/pyg-team/pytorch_geometric/pull/7421))
 - Added the `IGMCDataset` ([#7441](https://github.com/pyg-team/pytorch_geometric/pull/7441))
diff --git a/torch_geometric/distributed/dist_neighbor_sampler.py b/torch_geometric/distributed/dist_neighbor_sampler.py
index 1bdb33002..a6d7691c1 100644
--- a/torch_geometric/distributed/dist_neighbor_sampler.py
+++ b/torch_geometric/distributed/dist_neighbor_sampler.py
@@ -1,21 +1,49 @@
+import itertools
 import logging
-from typing import Any, Dict, List, Optional, Tuple, Union
 
+import numpy as np
+import torch
 import torch.multiprocessing as mp
+from torch import Tensor
 
 from torch_geometric.distributed import LocalFeatureStore, LocalGraphStore
 from torch_geometric.distributed.dist_context import DistContext, DistRole
-from torch_geometric.distributed.event_loop import ConcurrentEventLoop
+from torch_geometric.distributed.event_loop import (
+    ConcurrentEventLoop,
+    wrap_torch_future,
+)
 from torch_geometric.distributed.rpc import (
     RPCCallBase,
     RPCRouter,
+    rpc_async,
     rpc_partition_to_workers,
     rpc_register,
     shutdown_rpc,
 )
-from torch_geometric.sampler import NeighborSampler
+from torch_geometric.distributed.utils import (
+    BatchDict,
+    NodeDict,
+    remove_duplicates,
+)
+from torch_geometric.sampler import (
+    EdgeSamplerInput,
+    HeteroSamplerOutput,
+    NeighborSampler,
+    NodeSamplerInput,
+    SamplerOutput,
+)
 from torch_geometric.sampler.base import NumNeighbors, SubgraphType
-from torch_geometric.typing import EdgeType
+from torch_geometric.sampler.utils import remap_keys
+from torch_geometric.typing import (
+    Any,
+    Dict,
+    EdgeType,
+    List,
+    NodeType,
+    Optional,
+    Tuple,
+    Union,
+)
 
 NumNeighborsType = Union[NumNeighbors, List[int], Dict[EdgeType, List[int]]]
 
@@ -70,7 +98,8 @@ class DistNeighborSampler:
         self.temporal_strategy = temporal_strategy
         self.time_attr = time_attr
         self.with_edge_attr = self.dist_feature.has_edge_attr()
-        self.edge_permutation = None  # TODO: Debug edge_perm for LinkLoader
+        _, _, self.edge_permutation = self.dist_graph.csc()
+        self.csc = True
 
     def register_sampler_rpc(self) -> None:
         partition2workers = rpc_partition_to_workers(
@@ -97,6 +126,543 @@ class DistNeighborSampler:
         self.event_loop = ConcurrentEventLoop(self.concurrency)
         self.event_loop.start_loop()
 
+    # Node-based distributed sampling #########################################
+
+    def sample_from_nodes(
+            self, inputs: NodeSamplerInput,
+            **kwargs) -> Optional[Union[SamplerOutput, HeteroSamplerOutput]]:
+        inputs = NodeSamplerInput.cast(inputs)
+        if self.channel is None:
+            # synchronous sampling
+            return self.event_loop.run_task(
+                coro=self._sample_from(self.node_sample, inputs))
+
+        # asynchronous sampling
+        cb = kwargs.get("callback", None)
+        self.event_loop.add_task(
+            coro=self._sample_from(self.node_sample, inputs), callback=cb)
+        return None
+
+    async def _sample_from(
+        self,
+        async_func,
+        *args,
+        **kwargs,
+    ) -> Optional[Union[SamplerOutput, HeteroSamplerOutput]]:
+
+        sampler_output = await async_func(*args, **kwargs)
+        res = await self._collate_fn(sampler_output)
+
+        if self.channel is None:
+            return res
+        self.channel.put(res)
+        return None
+
+    async def node_sample(
+        self,
+        inputs: Union[NodeSamplerInput, EdgeSamplerInput],
+        dst_time: Tensor = None,
+    ) -> Union[SamplerOutput, HeteroSamplerOutput]:
+        r"""Performs layer by layer distributed sampling from a
+        :class:`NodeSamplerInput` and returns the output of the sampling
+        procedure.
+
+        Note:
+            In case of distributed training it is required to synchronize the
+            results between machines after each layer.
+        """
+
+        input_type = inputs.input_type
+        self.input_type = input_type
+        batch_size = inputs.input_id.size()[0]
+
+        seed_dict = None
+        seed_time_dict = None
+        src_batch_dict = None
+
+        if isinstance(inputs, NodeSamplerInput):
+            seed = inputs.node.to(self.device)
+            seed_time = (inputs.time.to(self.device)
+                         if inputs.time is not None else None)
+            src_batch = torch.arange(batch_size) if self.disjoint else None
+            seed_dict = {input_type: seed}
+            seed_time_dict: Dict[NodeType, Tensor] = {input_type: seed_time}
+            metadata = (seed, seed_time)
+
+        elif isinstance(inputs, EdgeSamplerInput) and self.is_hetero:
+            seed_dict = {input_type[0]: inputs.row, input_type[-1]: inputs.col}
+            if dst_time is not None:
+                seed_time_dict = {
+                    input_type[0]: inputs.time,
+                    input_type[-1]: dst_time,
+                }
+
+            if self.disjoint:
+                src_batch_dict = {
+                    input_type[0]: torch.arange(batch_size),
+                    input_type[-1]: torch.arange(batch_size, batch_size * 2),
+                }
+            metadata = (seed_dict, seed_time_dict)
+
+        else:
+            raise NotImplementedError
+
+        if self.is_hetero:
+            if input_type is None:
+                raise ValueError("Input type should be defined")
+
+            node_dict = NodeDict()
+            batch_dict = BatchDict(self.disjoint)
+            edge_dict: Dict[EdgeType, Tensor] = {}
+            num_sampled_nodes_dict: Dict[NodeType, List[int]] = {}
+            sampled_nbrs_per_node_dict: Dict[EdgeType, List[int]] = {}
+            num_sampled_edges_dict: Dict[EdgeType, List[int]] = {}
+
+            for ntype in self._sampler.node_types:
+                num_sampled_nodes_dict.update({ntype: [0]})
+
+            for etype in self._sampler.edge_types:
+                edge_dict.update({etype: torch.empty(0, dtype=torch.int64)})
+                num_sampled_edges_dict.update({etype: []})
+                sampled_nbrs_per_node_dict.update({etype: []})
+
+            if isinstance(inputs, EdgeSamplerInput):
+                node_dict.src = seed_dict
+                node_dict.out = {
+                    input_type[0]: inputs.row.numpy(),
+                    input_type[-1]: inputs.col.numpy(),
+                }
+
+                num_sampled_nodes_dict = {
+                    input_type[0]: [inputs.row.numel()],
+                    input_type[-1]: [inputs.col.numel()],
+                }
+
+                if self.disjoint:
+                    batch_dict = src_batch_dict
+                    batch_dict.out = {
+                        input_type[0]: src_batch_dict[input_type[0]].numpy(),
+                        input_type[-1]: src_batch_dict[input_type[-1]].numpy(),
+                    }
+
+            else:
+                node_dict.src[input_type] = seed
+                node_dict.out[input_type] = seed.numpy()
+
+                num_sampled_nodes_dict[input_type].append(seed.numel())
+
+                if self.disjoint:
+                    batch_dict.src[input_type] = src_batch
+                    batch_dict.out[input_type] = src_batch.numpy()
+
+            # loop over the layers
+            for i in range(self._sampler.num_hops):
+                # create tasks
+                task_dict = {}
+                for etype in self._sampler.edge_types:
+                    src = etype[0] if not self.csc else etype[2]
+
+                    if node_dict.src[src].numel():
+                        seed_time = (seed_time_dict.get(src, None)
+                                     if seed_time_dict is not None else None)
+                        if isinstance(self.num_neighbors, list):
+                            one_hop_num = self.num_neighbors[i]
+                        else:
+                            one_hop_num = self.num_neighbors[etype][i]
+
+                        task_dict[etype] = self.event_loop._loop.create_task(
+                            self.sample_one_hop(
+                                node_dict.src[src],
+                                one_hop_num,
+                                seed_time,
+                                batch_dict.src[src],
+                                etype,
+                            ))
+
+                for etype, task in task_dict.items():
+                    out: HeteroSamplerOutput = await task
+
+                    if out.node.numel() == 0:
+                        # no neighbors were sampled
+                        break
+
+                    dst = etype[2] if not self.csc else etype[0]
+
+                    # remove duplicates
+                    (
+                        node_dict.src[dst],
+                        node_dict.out[dst],
+                        batch_dict.src[dst],
+                        batch_dict.out[dst],
+                    ) = remove_duplicates(
+                        out,
+                        node_dict.out[dst],
+                        batch_dict.out[dst],
+                        self.disjoint,
+                    )
+
+                    node_dict.with_dupl[dst] = torch.cat(
+                        [node_dict.with_dupl[dst], out.node])
+                    edge_dict[etype] = torch.cat([edge_dict[etype], out.edge])
+
+                    if self.disjoint:
+                        batch_dict.with_dupl[dst] = torch.cat(
+                            [batch_dict.with_dupl[dst], out.batch])
+
+                    num_sampled_nodes_dict[dst].append(len(node_dict.src[dst]))
+                    num_sampled_edges_dict[etype].append(len(out.node))
+                    sampled_nbrs_per_node_dict[etype] += out.metadata
+
+            sampled_nbrs_per_node_dict = remap_keys(sampled_nbrs_per_node_dict,
+                                                    self._sampler.to_rel_type)
+
+            row_dict, col_dict = torch.ops.pyg.hetero_relabel_neighborhood(
+                self._sampler.node_types,
+                self._sampler.edge_types,
+                seed_dict,
+                node_dict.with_dupl,
+                sampled_nbrs_per_node_dict,
+                self._sampler.num_nodes,
+                batch_dict.with_dupl,
+                self.csc,
+                self.disjoint,
+            )
+
+            node_dict.out = {
+                ntype: torch.from_numpy(node_dict.out[ntype])
+                for ntype in self._sampler.node_types
+            }
+            if self.disjoint:
+                batch_dict.out = {
+                    ntype: torch.from_numpy(batch_dict.out[ntype])
+                    for ntype in self._sampler.node_types
+                }
+
+            sampler_output = HeteroSamplerOutput(
+                node=node_dict.out,
+                row=remap_keys(row_dict, self._sampler.to_edge_type),
+                col=remap_keys(col_dict, self._sampler.to_edge_type),
+                edge=edge_dict,
+                batch=batch_dict.out if self.disjoint else None,
+                num_sampled_nodes=num_sampled_nodes_dict,
+                num_sampled_edges=num_sampled_edges_dict,
+                metadata=metadata,
+            )
+        else:
+            src = seed
+            node = src.numpy()
+            batch = src_batch.numpy() if self.disjoint else None
+
+            node_with_dupl = [torch.empty(0, dtype=torch.int64)]
+            batch_with_dupl = [torch.empty(0, dtype=torch.int64)]
+            edge = [torch.empty(0, dtype=torch.int64)]
+
+            sampled_nbrs_per_node = []
+            num_sampled_nodes = [seed.numel()]
+            num_sampled_edges = [0]
+
+            # loop over the layers
+            for one_hop_num in self.num_neighbors:
+                out = await self.sample_one_hop(src, one_hop_num, seed_time,
+                                                src_batch)
+                if out.node.numel() == 0:
+                    # no neighbors were sampled
+                    break
+
+                # remove duplicates
+                src, node, src_batch, batch = remove_duplicates(
+                    out, node, batch, self.disjoint)
+
+                node_with_dupl.append(out.node)
+                edge.append(out.edge)
+
+                if self.disjoint:
+                    batch_with_dupl.append(out.batch)
+
+                num_sampled_nodes.append(len(src))
+                num_sampled_edges.append(len(out.node))
+                sampled_nbrs_per_node += out.metadata
+
+            row, col = torch.ops.pyg.relabel_neighborhood(
+                seed,
+                torch.cat(node_with_dupl),
+                sampled_nbrs_per_node,
+                self._sampler.num_nodes,
+                torch.cat(batch_with_dupl) if self.disjoint else None,
+                self.csc,
+                self.disjoint,
+            )
+
+            sampler_output = SamplerOutput(
+                node=torch.from_numpy(node),
+                row=row,
+                col=col,
+                edge=torch.cat(edge),
+                batch=torch.from_numpy(batch) if self.disjoint else None,
+                num_sampled_nodes=num_sampled_nodes,
+                num_sampled_edges=num_sampled_edges,
+                metadata=metadata,
+            )
+
+        if self.subgraph_type == SubgraphType.bidirectional:
+            sampler_output = sampler_output.to_bidirectional()
+
+        return sampler_output
+
+    def get_sampler_output(
+        self,
+        outputs: List[SamplerOutput],
+        seed_size: int,
+        p_id: int,
+        src_batch: Optional[Tensor] = None,
+    ) -> SamplerOutput:
+        r"""Used when seed nodes belongs to one partition. It's purpose is to
+        remove seed nodes from sampled nodes and calculates how many neighbors
+        were sampled by each src node based on the
+        :obj:`cumsum_neighbors_per_node`. Returns updated sampler output.
+        """
+        cumsum_neighbors_per_node = outputs[p_id].metadata
+
+        # do not include seed
+        outputs[p_id].node = outputs[p_id].node[seed_size:]
+
+        begin = np.array(cumsum_neighbors_per_node[1:])
+        end = np.array(cumsum_neighbors_per_node[:-1])
+
+        sampled_nbrs_per_node = list(np.subtract(begin, end))
+
+        outputs[p_id].metadata = sampled_nbrs_per_node
+
+        if self.disjoint:
+            batch = [[src_batch[i]] * nbrs_per_node
+                     for i, nbrs_per_node in enumerate(sampled_nbrs_per_node)]
+            outputs[p_id].batch = Tensor(
+                list(itertools.chain.from_iterable(batch))).type(torch.int64)
+
+        return outputs[p_id]
+
+    def merge_sampler_outputs(
+        self,
+        partition_ids: Tensor,
+        partition_orders: Tensor,
+        outputs: List[SamplerOutput],
+        one_hop_num: int,
+        src_batch: Optional[Tensor] = None,
+    ) -> SamplerOutput:
+        r"""Merges samplers outputs from different partitions, so that they
+        are sorted according to the sampling order. Removes seed nodes from
+        sampled nodes and calculates how many neighbors were sampled by each
+        src node based on the :obj:`cumsum_neighbors_per_node`. Leverages the
+        :obj:`pyg-lib` :obj:`merge_sampler_outputs` function.
+
+        Args:
+            partition_ids (torch.Tensor): Contains information on which
+                partition seeds nodes are located on.
+            partition_orders (torch.Tensor): Contains information about the
+                order of seed nodes in each partition.
+            outputs (List[SamplerOutput]): List of all samplers outputs.
+            one_hop_num (int): Max number of neighbors sampled in the current
+                layer.
+
+        Returns :obj:`SamplerOutput` containing all merged outputs.
+        """
+        sampled_nodes_with_dupl = [
+            o.node if o is not None else torch.empty(0, dtype=torch.int64)
+            for o in outputs
+        ]
+        edge_ids = [
+            o.edge if o is not None else torch.empty(0, dtype=torch.int64)
+            for o in outputs
+        ]
+        cumm_sampled_nbrs_per_node = [
+            o.metadata if o is not None else [] for o in outputs
+        ]
+
+        partition_ids = partition_ids.tolist()
+        partition_orders = partition_orders.tolist()
+
+        partitions_num = self.dist_graph.meta["num_parts"]
+
+        out = torch.ops.pyg.merge_sampler_outputs(
+            sampled_nodes_with_dupl,
+            edge_ids,
+            cumm_sampled_nbrs_per_node,
+            partition_ids,
+            partition_orders,
+            partitions_num,
+            one_hop_num,
+            src_batch,
+            self.disjoint,
+        )
+        (
+            out_node_with_dupl,
+            out_edge,
+            out_batch,
+            out_sampled_nbrs_per_node,
+        ) = out
+
+        return SamplerOutput(
+            out_node_with_dupl,
+            None,
+            None,
+            out_edge,
+            out_batch if self.disjoint else None,
+            metadata=(out_sampled_nbrs_per_node),
+        )
+
+    async def sample_one_hop(
+        self,
+        srcs: Tensor,
+        one_hop_num: int,
+        seed_time: Optional[Tensor] = None,
+        src_batch: Optional[Tensor] = None,
+        etype: Optional[EdgeType] = None,
+    ) -> SamplerOutput:
+        r"""Sample one-hop neighbors for a :obj:`srcs`. If src node is located
+        on a local partition, evaluates the :obj:`_sample_one_hop` function on
+        a current machine. If src node is from a remote partition, send a
+        request to a remote machine that contains this partition.
+
+        Returns merged samplers outputs from local / remote machines.
+        """
+        partition_ids = self.dist_graph.get_partition_ids_from_nids(srcs)
+        partition_orders = torch.zeros(len(partition_ids), dtype=torch.long)
+
+        p_outputs: List[SamplerOutput] = [None
+                                          ] * self.dist_graph.meta["num_parts"]
+        futs: List[torch.futures.Future] = []
+
+        local_only = True
+        single_partition = len(set(partition_ids.tolist())) == 1
+
+        for i in range(self.dist_graph.num_partitions):
+            p_id = (self.dist_graph.partition_idx +
+                    i) % self.dist_graph.num_partitions
+            p_mask = partition_ids == p_id
+            p_srcs = torch.masked_select(srcs, p_mask)
+            p_seed_time = (torch.masked_select(seed_time, p_mask)
+                           if seed_time is not None else None)
+
+            p_indices = torch.arange(len(p_srcs), dtype=torch.long)
+            partition_orders[p_mask] = p_indices
+
+            if p_srcs.shape[0] > 0:
+                if p_id == self.dist_graph.partition_idx:
+                    # sample on local machine
+                    p_nbr_out = self._sampler._sample_one_hop(
+                        p_srcs, one_hop_num, p_seed_time, self.csc, etype)
+                    p_outputs.pop(p_id)
+                    p_outputs.insert(p_id, p_nbr_out)
+                else:
+                    # sample on remote machine
+                    local_only = False
+                    to_worker = self.rpc_router.get_to_worker(p_id)
+                    futs.append(
+                        rpc_async(
+                            to_worker,
+                            self.rpc_sample_callee_id,
+                            args=(
+                                p_srcs,
+                                one_hop_num,
+                                p_seed_time,
+                                self.csc,
+                                etype,
+                            ),
+                        ))
+
+        if not local_only:
+            # Src nodes are remote
+            res_fut_list = await wrap_torch_future(
+                torch.futures.collect_all(futs))
+            for i, res_fut in enumerate(res_fut_list):
+                p_id = (self.dist_graph.partition_idx + i +
+                        1) % self.dist_graph.num_partitions
+                p_outputs.pop(p_id)
+                p_outputs.insert(p_id, res_fut.wait())
+
+        # All src nodes are in the same partition
+        if single_partition:
+            return self.get_sampler_output(p_outputs, len(srcs),
+                                           partition_ids[0], src_batch)
+
+        return self.merge_sampler_outputs(partition_ids, partition_orders,
+                                          p_outputs, one_hop_num, src_batch)
+
+    async def _collate_fn(
+        self, output: Union[SamplerOutput, HeteroSamplerOutput]
+    ) -> Union[SamplerOutput, HeteroSamplerOutput]:
+        r"""Collect labels and features for the sampled subgrarph if necessary,
+        and put them into a sample message.
+        """
+        if self.is_hetero:
+            nlabels = {}
+            nfeats = {}
+            efeats = {}
+            # Collect node labels of input node type.
+            node_labels = self.dist_feature.labels
+            if node_labels is not None:
+                nlabels = node_labels[output.node[self.input_type]]
+            else:
+                nlabels = None
+            # Collect node features.
+            if output.node is not None:
+                for ntype in output.node.keys():
+                    if output.node[ntype].numel() > 0:
+                        fut = self.dist_feature.lookup_features(
+                            is_node_feat=True,
+                            index=output.node[ntype],
+                            input_type=ntype,
+                        )
+                        nfeat = await wrap_torch_future(fut)
+                        nfeat = nfeat.to(torch.device("cpu"))
+                        nfeats[ntype] = nfeat
+                    else:
+                        nfeats[ntype] = None
+            # Collect edge features
+            if output.edge is not None and self.with_edge_attr:
+                for etype in output.edge.keys():
+                    if output.edge[etype].numel() > 0:
+                        fut = self.dist_feature.lookup_features(
+                            is_node_feat=False,
+                            index=output.edge[etype],
+                            input_type=etype,
+                        )
+                        efeat = await wrap_torch_future(fut)
+                        efeat = efeat.to(torch.device("cpu"))
+                        efeats[etype] = efeat
+                    else:
+                        efeats[etype] = None
+
+        else:  # Homo
+            # Collect node labels.
+            nlabels = (self.dist_feature.labels[output.node] if
+                       (self.dist_feature.labels is not None) else None)
+            # Collect node features.
+            if output.node is not None:
+                fut = self.dist_feature.lookup_features(
+                    is_node_feat=True, index=output.node)
+                nfeats = await wrap_torch_future(fut)
+                nfeats = nfeats.to(torch.device("cpu"))
+            # else:
+            efeats = None
+            # Collect edge features.
+            if output.edge is not None and self.with_edge_attr:
+                fut = self.dist_feature.lookup_features(
+                    is_node_feat=False, index=output.edge)
+                efeats = await wrap_torch_future(fut)
+                efeats = efeats.to(torch.device("cpu"))
+            else:
+                efeats = None
+
+        output.metadata = (*output.metadata, nfeats, nlabels, efeats)
+        if self.is_hetero:
+            output.row = remap_keys(output.row, self._sampler.to_edge_type)
+            output.col = remap_keys(output.col, self._sampler.to_edge_type)
+        return output
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()-PID{mp.current_process().pid}"
+
 
 # Sampling Utilities ##########################################################
 
diff --git a/torch_geometric/distributed/event_loop.py b/torch_geometric/distributed/event_loop.py
new file mode 100644
index 000000000..0d766d245
--- /dev/null
+++ b/torch_geometric/distributed/event_loop.py
@@ -0,0 +1,97 @@
+import asyncio
+import logging
+from threading import BoundedSemaphore, Thread
+from typing import Callable, Optional
+
+import torch
+
+# This code is taken from alibaba graphlearn-for-pytorch repository:
+# https://github.com/alibaba/graphlearn-for-pytorch/blob/main/graphlearn_torch/
+# python/distributed/event_loop.py
+
+
+def to_asyncio_future(future: torch.futures.Future) -> asyncio.futures.Future:
+    r"""Convert a :class:`torch.futures.Future` to a standard asyncio future.
+    """
+    loop = asyncio.get_event_loop()
+    asyncio_future = loop.create_future()
+
+    def on_done(*_):
+        try:
+            result = future.wait()
+        except Exception as e:
+            loop.call_soon_threadsafe(asyncio_future.set_exception, e)
+        else:
+            loop.call_soon_threadsafe(asyncio_future.set_result, result)
+
+    future.add_done_callback(on_done)
+
+    return asyncio_future
+
+
+class ConcurrentEventLoop:
+    r"""Concurrent event loop context.
+
+    Args:
+        concurrency: max processing concurrency.
+    """
+    def __init__(self, concurrency: int):
+        self._concurrency = concurrency
+        self._sem = BoundedSemaphore(concurrency)
+        self._loop = asyncio.new_event_loop()
+        self._runner_t = Thread(target=self._run_loop)
+        self._runner_t.daemon = True
+
+    def start_loop(self):
+        if not self._runner_t.is_alive():
+            self._runner_t.start()
+
+    def shutdown_loop(self):
+        self.wait_all()
+        if self._runner_t.is_alive():
+            self._loop.stop()
+            self._runner_t.join(timeout=1)
+
+    def wait_all(self):
+        r"""Wait for all pending tasks to be finished."""
+        for _ in range(self._concurrency):
+            self._sem.acquire()
+        for _ in range(self._concurrency):
+            self._sem.release()
+
+    def add_task(self, coro, callback: Optional[Callable] = None):
+        r"""Adds an asynchronized coroutine task to run.
+
+        Args:
+            coro: The asynchronous coroutine function.
+            callback (callable, optional): The callback function applied on the
+                returned results after the coroutine task is finished.
+                (default: :obj:`None`)
+
+        Note that any result returned by :obj:`callback` will be ignored.
+        """
+        def on_done(f: asyncio.futures.Future):
+            try:
+                res = f.result()
+                if callback is not None:
+                    callback(res)
+            except Exception as e:
+                logging.error(f"Coroutine task failed with error: {e}")
+            self._sem.release()
+
+        self._sem.acquire()
+        fut = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        fut.add_done_callback(on_done)
+
+    def run_task(self, coro):
+        r"""Runs a coroutine task synchronously.
+
+        Args:
+            coro: The synchronous coroutine function.
+        """
+        with self._sem:
+            fut = asyncio.run_coroutine_threadsafe(coro, self._loop)
+            return fut.result()
+
+    def _run_loop(self):
+        self._loop.run_forever()
diff --git a/torch_geometric/distributed/utils.py b/torch_geometric/distributed/utils.py
index 0308a5f99..ea3917507 100644
--- a/torch_geometric/distributed/utils.py
+++ b/torch_geometric/distributed/utils.py
@@ -1,10 +1,74 @@
-from typing import Dict, Optional
+from dataclasses import dataclass, field
+from typing import Tuple
 
+import numpy as np
 import torch
 from torch import Tensor
 
 from torch_geometric.data import HeteroData
 from torch_geometric.distributed import LocalFeatureStore, LocalGraphStore
+from torch_geometric.sampler import SamplerOutput
+from torch_geometric.typing import Dict, NodeType, Optional
+
+
+@dataclass
+class NodeDict:
+    r"""Class used during heterogeneous sampling:
+    1) The nodes to serve as source nodes in the next layer
+    2) The nodes with duplicates that are further needed to create COO output
+    3) The output nodes without duplicates
+    """
+    src: Dict[NodeType, Tensor] = field(default_factory=dict)
+    with_dupl: Dict[NodeType, Tensor] = field(default_factory=dict)
+    out: Dict[NodeType, Tensor] = field(default_factory=dict)
+
+
+@dataclass
+class BatchDict:
+    r"""Class used during disjoint heterogeneous sampling:
+    1) The batch to serve as initial subgraph IDs for source nodes in the next
+       layer
+    2) The subgraph IDs with duplicates that are further needed to create COO
+       output
+    3) The output subgraph IDs without duplicates
+    """
+    src: Dict[NodeType, Tensor] = field(default_factory=dict)
+    with_duple: Dict[NodeType, Tensor] = field(default_factory=dict)
+    out: Dict[NodeType, Tensor] = field(default_factory=dict)
+
+
+def remove_duplicates(
+    out: SamplerOutput,
+    node: Tensor,
+    batch: Optional[Tensor] = None,
+    disjoint: bool = False,
+) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:
+
+    num_nodes = node.numel()
+    node_combined = torch.cat([node, out.node])
+
+    if not disjoint:
+        _, idx = np.unique(node_combined.cpu().numpy(), return_index=True)
+        idx = torch.from_numpy(idx).to(node.device).sort().values
+
+        node = node_combined[idx]
+        src = node[num_nodes:]
+
+        return (src, node, None, None)
+
+    else:
+        batch_combined = torch.cat([batch, out.batch])
+        node_batch = torch.stack([batch_combined, node_combined], dim=0)
+
+        _, idx = np.unique(node_batch.cpu().numpy(), axis=1, return_index=True)
+        idx = torch.from_numpy(idx).to(node.device).sort().values
+
+        batch = batch_combined[idx]
+        node = node_combined[idx]
+        src_batch = batch[num_nodes:]
+        src = node[num_nodes:]
+
+        return (src, node, src_batch, batch)
 
 
 def filter_dist_store(
diff --git a/torch_geometric/sampler/neighbor_sampler.py b/torch_geometric/sampler/neighbor_sampler.py
index 2a4c3d0aa..0d68a52b9 100644
--- a/torch_geometric/sampler/neighbor_sampler.py
+++ b/torch_geometric/sampler/neighbor_sampler.py
@@ -124,12 +124,12 @@ class NeighborSampler(BaseSampler):
             edge_attrs = graph_store.get_all_edge_attrs()
             self.edge_types = list(set(attr.edge_type for attr in edge_attrs))
 
-            self.num_nodes = {
-                node_type: remote_backend_utils.size(*data, node_type)
-                for node_type in self.node_types
-            }
+            if weight_attr is not None:
+                raise NotImplementedError(
+                    f"'weight_attr' argument not yet supported within "
+                    f"'{self.__class__.__name__}' for "
+                    f"'(FeatureStore, GraphStore)' inputs")
 
-            self.node_time: Optional[Dict[NodeType, Tensor]] = None
             if time_attr is not None:
                 # If the `time_attr` is present, we expect that `GraphStore`
                 # holds all edges sorted by destination, and within local
@@ -151,29 +151,47 @@ class NeighborSampler(BaseSampler):
                     copy.copy(attr) for attr in node_attrs
                     if attr.attr_name == time_attr
                 ]
-                for attr in time_attrs:  # Reset the index to obtain full data.
-                    attr.index = None
-                time_tensors = feature_store.multi_get_tensor(time_attrs)
-                self.node_time = {
-                    time_attr.group_name: time_tensor
-                    for time_attr, time_tensor in zip(time_attrs, time_tensors)
-                }
 
-            self.edge_weight: Optional[Dict[EdgeType, Tensor]] = None
-            if weight_attr is not None:
-                raise NotImplementedError(
-                    f"'weight_attr' argument not yet supported within "
-                    f"'{self.__class__.__name__}' for "
-                    f"'(FeatureStore, GraphStore)' inputs")
+            if not self.is_hetero:
+                self.num_nodes = max(edge_attrs[0].size)
+                self.edge_weight: Optional[Tensor] = None
 
-            # Conversion to/from C++ string type (see above):
-            self.to_rel_type = {k: '__'.join(k) for k in self.edge_types}
-            self.to_edge_type = {v: k for k, v in self.to_rel_type.items()}
+                self.node_time: Optional[Tensor] = None
+                if time_attr is not None and len(time_attrs) != 1:
+                    raise ValueError("Temporal sampling specified but did "
+                                     "not find any temporal data")
 
-            # Convert the graph data into CSC format for sampling:
-            row_dict, colptr_dict, self.perm = graph_store.csc()
-            self.row_dict = remap_keys(row_dict, self.to_rel_type)
-            self.colptr_dict = remap_keys(colptr_dict, self.to_rel_type)
+                    time_attrs[0].index = None  # Reset index for full data.
+                    time_tensor = feature_store.get_tensor(time_attrs[0])
+                    self.node_time = time_tensor
+
+                self.row, self.colptr, self.perm = graph_store.csc()
+
+            else:
+                self.num_nodes = {
+                    node_type: remote_backend_utils.size(*data, node_type)
+                    for node_type in self.node_types
+                }
+                self.edge_weight: Optional[Dict[EdgeType, Tensor]] = None
+
+                self.node_time: Optional[Dict[NodeType, Tensor]] = None
+                if time_attr is not None:
+                    for attr in time_attrs:  # Reset index for full data.
+                        attr.index = None
+                    time_tensors = feature_store.multi_get_tensor(time_attrs)
+                    self.node_time = {
+                        time_attr.group_name: time_tensor
+                        for time_attr, time_tensor in zip(
+                            time_attrs, time_tensors)
+                    }
+
+                # Conversion to/from C++ string type (see above):
+                self.to_rel_type = {k: '__'.join(k) for k in self.edge_types}
+                self.to_edge_type = {v: k for k, v in self.to_rel_type.items()}
+                # Convert the graph data into CSC format for sampling:
+                row_dict, colptr_dict, self.perm = graph_store.csc()
+                self.row_dict = remap_keys(row_dict, self.to_rel_type)
+                self.colptr_dict = remap_keys(colptr_dict, self.to_rel_type)
 
         if (self.edge_weight is not None
                 and not torch_geometric.typing.WITH_WEIGHTED_NEIGHBOR_SAMPLE):
@@ -197,6 +215,16 @@ class NeighborSampler(BaseSampler):
         else:
             self._num_neighbors = NumNeighbors(num_neighbors)
 
+    @property
+    def is_hetero(self) -> bool:
+        if self.data_type == DataType.homogeneous:
+            return False
+        if self.data_type == DataType.heterogeneous:
+            return True
+
+        # self.data_type == DataType.remote
+        return self.node_types != [None]
+
     @property
     def is_temporal(self) -> bool:
         return self.node_time is not None
@@ -223,8 +251,9 @@ class NeighborSampler(BaseSampler):
     # Edge-based sampling #####################################################
 
     def sample_from_edges(
-        self, inputs: EdgeSamplerInput,
-        neg_sampling: Optional[NegativeSampling] = None
+        self,
+        inputs: EdgeSamplerInput,
+        neg_sampling: Optional[NegativeSampling] = None,
     ) -> Union[SamplerOutput, HeteroSamplerOutput]:
         out = edge_sample(inputs, self._sample, self.num_nodes, self.disjoint,
                           self.node_time, neg_sampling)
@@ -284,8 +313,8 @@ class NeighborSampler(BaseSampler):
 
                 # `pyg-lib>0.1.0` returns sampled number of nodes/edges:
                 num_sampled_nodes = num_sampled_edges = None
-                if len(out) == 6:
-                    num_sampled_nodes, num_sampled_edges = out[4:]
+                if len(out) >= 6:
+                    num_sampled_nodes, num_sampled_edges = out[4:6]
 
                 if self.disjoint:
                     node = {k: v.t().contiguous() for k, v in node.items()}
@@ -364,8 +393,8 @@ class NeighborSampler(BaseSampler):
 
                 # `pyg-lib>0.1.0` returns sampled number of nodes/edges:
                 num_sampled_nodes = num_sampled_edges = None
-                if len(out) == 6:
-                    num_sampled_nodes, num_sampled_edges = out[4:]
+                if len(out) >= 6:
+                    num_sampled_nodes, num_sampled_edges = out[4:6]
 
                 if self.disjoint:
                     batch, node = node.t().contiguous()
@@ -402,6 +431,56 @@ class NeighborSampler(BaseSampler):
                 num_sampled_edges=num_sampled_edges,
             )
 
+    def _sample_one_hop(
+        self,
+        input_nodes: Tensor,
+        num_neighbors: int,
+        seed_time: Optional[Tensor] = None,
+        edge_type: Optional[EdgeType] = None,
+    ) -> SamplerOutput:
+        r"""Implements one-hop neighbor sampling for a set of input nodes for a
+        specific edge type.
+        """
+        rel_type = '__'.join(edge_type) if self.is_hetero else None
+
+        if not self.is_hetero:
+            colptr = self.colptr
+            row = self.row
+            node_time = self.node_time
+        else:
+            rel_type = '__'.join(edge_type)
+            colptr = self.colptr_dict[rel_type]
+            row = self.row_dict[rel_type]
+            node_time = self.node_time.get(edge_type[2], None)
+
+        out = torch.ops.pyg.dist_neighbor_sample(
+            colptr,
+            row,
+            input_nodes.to(colptr.dtype),
+            num_neighbors,
+            node_time,
+            seed_time,
+            None,  # TODO: edge_weight
+            True,  # csc
+            self.replace,
+            self.subgraph_type != SubgraphType.induced,
+            self.disjoint and node_time is not None,
+            self.temporal_strategy,
+        )
+        node, edge, cumsum_neighbors_per_node = out
+
+        if self.disjoint:
+            batch, node = node.t().contiguous()
+
+        return SamplerOutput(
+            node=node,
+            row=None,
+            col=None,
+            edge=edge,
+            batch=batch,
+            metadata=(cumsum_neighbors_per_node, ),
+        )
+
 
 # Sampling Utilities ##########################################################
 
@@ -598,8 +677,13 @@ def edge_sample(
 
             dst_neg_index = dst_neg_index.view(num_pos, -1).squeeze(-1)
 
-            out.metadata = (input_id, src_index, dst_pos_index, dst_neg_index,
-                            src_time)
+            out.metadata = (
+                input_id,
+                src_index,
+                dst_pos_index,
+                dst_neg_index,
+                src_time,
+            )
 
     # Homogeneus Neighborhood Sampling ########################################
 
@@ -641,15 +725,24 @@ def edge_sample(
                 dst_neg_index = inverse_seed[2 * num_pos:]
             dst_neg_index = dst_neg_index.view(num_pos, -1).squeeze(-1)
 
-            out.metadata = (input_id, src_index, dst_pos_index, dst_neg_index,
-                            src_time)
+            out.metadata = (
+                input_id,
+                src_index,
+                dst_pos_index,
+                dst_neg_index,
+                src_time,
+            )
 
     return out
 
 
-def neg_sample(seed: Tensor, neg_sampling: NegativeSampling, num_nodes: int,
-               seed_time: Optional[Tensor],
-               node_time: Optional[Tensor]) -> Tensor:
+def neg_sample(
+    seed: Tensor,
+    neg_sampling: NegativeSampling,
+    num_nodes: int,
+    seed_time: Optional[Tensor],
+    node_time: Optional[Tensor],
+) -> Tensor:
     num_neg = math.ceil(seed.numel() * neg_sampling.amount)
 
     # TODO: Do not sample false negatives.


Test patch:
diff --git a/test/distributed/test_dist_utils.py b/test/distributed/test_dist_utils.py
new file mode 100644
index 000000000..d88bc94f2
--- /dev/null
+++ b/test/distributed/test_dist_utils.py
@@ -0,0 +1,35 @@
+import torch
+
+from torch_geometric.distributed.utils import remove_duplicates
+from torch_geometric.sampler import SamplerOutput
+
+
+def test_remove_duplicates():
+    node = torch.tensor([0, 1, 2, 3])
+    out_node = torch.tensor([0, 4, 1, 5, 1, 6, 2, 7, 3, 8])
+
+    out = SamplerOutput(out_node, None, None, None)
+
+    src, node, _, _ = remove_duplicates(out, node)
+
+    assert src.tolist() == [4, 5, 6, 7, 8]
+    assert node.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8]
+
+
+def test_remove_duplicates_disjoint():
+    node = torch.tensor([0, 1, 2, 3])
+    batch = torch.tensor([0, 1, 2, 3])
+
+    out_node = torch.tensor([0, 4, 1, 5, 1, 6, 2, 6, 7, 3, 8])
+    out_batch = torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3])
+
+    out = SamplerOutput(out_node, None, None, None, out_batch)
+
+    src, node, src_batch, batch = remove_duplicates(out, node, batch,
+                                                    disjoint=True)
+
+    assert src.tolist() == [4, 5, 6, 7, 8]
+    assert node.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8]
+
+    assert src_batch.tolist() == [0, 1, 2, 3, 3]
+    assert batch.tolist() == [0, 1, 2, 3, 0, 1, 2, 3, 3]