Repository: pyg-team/pytorch_geometric
Pull request #7671 (2023-06-30 22:17:51)

Add `dist_context`/RPC for distributed training
This code belongs to the part of the whole distributed training for PyG.

This class will do

1. dist_context is to setup the distributed mode like worker mode or other mode and also setup the information for distributed context like role, rank, group_name, world_size,etc.
2. based on pytorch rpc.api to setup the api to check is_rpc_initialized and wrapper func, init_rpc(), RpcCall(), RpcRouter, rpc_async_request() and rpc_sync_request(), etc

These basic rpc functionality will be used in later feature lookup function after node sampling. 

Any comments please let us know. thanks.



codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7671?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7671](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7671?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (e8b6cf8) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/125b59560c23a9dfa09bf931e4e93069b8745930?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (125b595) will **decrease** coverage by `0.49%`.
> The diff coverage is `0.00%`.

```diff
@@            Coverage Diff             @@
##           master    #7671      +/-   ##
==========================================
- Coverage   91.58%   91.09%   -0.49%     
==========================================
  Files         452      454       +2     
  Lines       25438    25639     +201     
==========================================
+ Hits        23297    23357      +60     
- Misses       2141     2282     +141     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7671?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/distributed/dist\_context.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7671?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL2Rpc3RfY29udGV4dC5weQ==) | `0.00% <0.00%> (ø)` | |
| [torch\_geometric/distributed/rpc.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7671?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL3JwYy5weQ==) | `0.00% <0.00%> (ø)` | |

... and [29 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7671/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 f0e9bdb14..a6ba8b339 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -33,7 +33,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))
 - 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))
+- 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))
 - 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_context.py b/torch_geometric/distributed/dist_context.py
new file mode 100644
index 000000000..5b3e72f73
--- /dev/null
+++ b/torch_geometric/distributed/dist_context.py
@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+from enum import Enum
+
+
+class DistRole(Enum):
+    WORKER = 1
+
+
+@dataclass
+class DistContext:
+    r"""Context information of the current process."""
+    rank: int
+    global_rank: int
+    world_size: int
+    global_world_size: int
+    group_name: str
+    role: DistRole = DistRole.WORKER
+
+    @property
+    def worker_name(self) -> str:
+        return f'{self.group_name}-{self.rank}'
diff --git a/torch_geometric/distributed/rpc.py b/torch_geometric/distributed/rpc.py
new file mode 100644
index 000000000..3f710d9b8
--- /dev/null
+++ b/torch_geometric/distributed/rpc.py
@@ -0,0 +1,190 @@
+import atexit
+import logging
+import threading
+from abc import ABC, abstractmethod
+from typing import Dict, List
+
+from torch._C._distributed_rpc import _is_current_rpc_agent_set
+from torch.distributed import rpc
+
+from torch_geometric.distributed.dist_context import DistContext, DistRole
+
+_rpc_init_lock = threading.RLock()
+
+
+def rpc_is_initialized() -> bool:
+    return _is_current_rpc_agent_set()
+
+
+@rpc.api._require_initialized
+def global_all_gather(obj, timeout=None):
+    r"""Gathers objects from all groups in a list."""
+    if timeout is None:
+        return rpc.api._all_gather(obj)
+    return rpc.api._all_gather(obj, timeout=timeout)
+
+
+@rpc.api._require_initialized
+def global_barrier(timeout=None):
+    r""" Block until all local and remote RPC processes."""
+    try:
+        global_all_gather(obj=None, timeout=timeout)
+    except RuntimeError:
+        logging.error("Failed to respond to global barrier")
+
+
+def init_rpc(
+    current_ctx: DistContext,
+    rpc_worker_names: Dict[DistRole, List[str]],
+    master_addr: str,
+    master_port: int,
+    num_rpc_threads: int = 16,
+    rpc_timeout: float = 240,
+):
+    with _rpc_init_lock:
+        if rpc_is_initialized():
+            return
+
+        if current_ctx is None:
+            raise RuntimeError("'dist_context' has not been set in 'init_rpc'")
+
+        options = rpc.TensorPipeRpcBackendOptions(
+            _transports=['ibv', 'uv'],
+            _channels=['mpt_uv', 'basic'],
+            num_worker_threads=num_rpc_threads,
+            rpc_timeout=rpc_timeout,
+            init_method=f'tcp://{master_addr}:{master_port}',
+        )
+
+        rpc.init_rpc(
+            name=current_ctx.worker_name,
+            rank=current_ctx.global_rank,
+            world_size=current_ctx.global_world_size,
+            rpc_backend_options=options,
+        )
+
+        gathered_results = global_all_gather(
+            obj=(current_ctx.role, current_ctx.world_size, current_ctx.rank),
+            timeout=rpc_timeout,
+        )
+
+        for worker_name, (role, world_size, rank) in gathered_results.items():
+            worker_list = rpc_worker_names.get(role, None)
+            if worker_list is None:
+                worker_list = [None for _ in range(world_size)]
+            else:
+                if len(worker_list) != world_size:
+                    raise RuntimeError(f"Inconsistent world size found in "
+                                       f"'init_rpc' (got {len(worker_list)})")
+
+            worker_list[rank] = worker_name
+            rpc_worker_names[role] = worker_list
+
+        global_barrier(timeout=rpc_timeout)
+
+
+def shutdown_rpc(graceful: bool = True):
+    if rpc_is_initialized():
+        rpc.shutdown(graceful)
+
+
+atexit.register(shutdown_rpc, False)
+
+
+class RpcRouter:
+    r"""A router to get the worker based on the partition ID."""
+    def __init__(self, partition_to_workers: List[List[str]]):
+        for pid, rpc_worker_list in enumerate(partition_to_workers):
+            if len(rpc_worker_list) == 0:
+                raise ValueError("No RPC worker is in worker list")
+        self.partition_to_workers = partition_to_workers
+        self.rpc_worker_indices = [0 for _ in range(len(partition_to_workers))]
+
+    def get_to_worker(self, partition_idx: int) -> str:
+        rpc_worker_list = self.partition_to_workers[partition_idx]
+        worker_idx = self.rpc_worker_indices[partition_idx]
+        router_worker = rpc_worker_list[worker_idx]
+        self.rpc_worker_indices[partition_idx] = ((worker_idx + 1) %
+                                                  len(rpc_worker_list))
+        return router_worker
+
+
+@rpc.api._require_initialized
+def rpc_partition_to_workers(
+    current_ctx: DistContext,
+    num_partitions: int,
+    current_partition_idx: int,
+):
+    r"""Performs an :obj:`all_gather` to get the mapping between partition and
+    workers."""
+    ctx = current_ctx
+    partition_to_workers = [[] for _ in range(num_partitions)]
+    gathered_results = global_all_gather(
+        (ctx.role, num_partitions, current_partition_idx))
+    for worker_name, (role, nparts, idx) in gathered_results.items():
+        partition_to_workers[idx].append(worker_name)
+    return partition_to_workers
+
+
+class RpcCallBase(ABC):
+    r"""A wrapper base class for RPC calls in remote processes."""
+    @abstractmethod
+    def rpc_sync(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def rpc_async(self, *args, **kwargs):
+        pass
+
+
+_rpc_call_lock = threading.RLock()
+_rpc_call_id: int = 0
+_rpc_call_pool: Dict[int, RpcCallBase] = {}
+
+
+@rpc.api._require_initialized
+def rpc_register(call: RpcCallBase) -> int:
+    r"""Registers a call for RPC requests."""
+    global _rpc_call_id, _rpc_call_pool
+
+    with _rpc_call_lock:
+        call_id = _rpc_call_id
+        _rpc_call_id += 1
+        if call_id in _rpc_call_pool:
+            raise RuntimeError("Registered function twice in 'rpc_register'")
+        _rpc_call_pool[call_id] = call
+
+    return call_id
+
+
+def _rpc_async_call(call_id: int, *args, **kwargs):
+    r""" Entry point for RPC requests."""
+    return _rpc_call_pool.get(call_id).rpc_async(*args, **kwargs)
+
+
+@rpc.api._require_initialized
+def rpc_async(worker_name: str, call_id: int, args=None, kwargs=None):
+    r"""Performs an asynchronous RPC request and returns a future."""
+    return rpc.rpc_async(
+        to=worker_name,
+        func=_rpc_async_call,
+        args=(call_id, *args),
+        kwargs=kwargs,
+    )
+
+
+def _rpc_sync_call(call_id: int, *args, **kwargs):
+    r"""Entry point for synchronous RPC requests."""
+    return _rpc_call_pool.get(call_id).rpc_sync(*args, **kwargs)
+
+
+@rpc.api._require_initialized
+def rpc_sync(worker_name: str, call_id: int, args=None, kwargs=None):
+    r"""Performs a synchronous RPC request and returns a future."""
+    future = rpc.rpc_async(
+        to=worker_name,
+        func=_rpc_sync_call,
+        args=(call_id, *args),
+        kwargs=kwargs,
+    )
+    return future.wait()


Test patch:
diff --git a/test/distributed/test_rpc.py b/test/distributed/test_rpc.py
new file mode 100644
index 000000000..673c53a69
--- /dev/null
+++ b/test/distributed/test_rpc.py
@@ -0,0 +1,121 @@
+import socket
+from typing import Dict, List
+
+import torch
+
+import torch_geometric.distributed.rpc as rpc
+from torch_geometric.distributed import LocalFeatureStore
+from torch_geometric.distributed.dist_context import DistContext, DistRole
+from torch_geometric.distributed.rpc import RpcRouter
+
+
+def run_rpc_feature_test(
+    world_size: int,
+    rank: int,
+    feature: LocalFeatureStore,
+    partition_book: torch.Tensor,
+    master_port: int,
+):
+    # 1) Initialize the context info:
+    current_ctx = DistContext(
+        rank=rank,
+        global_rank=rank,
+        world_size=world_size,
+        global_world_size=world_size,
+        group_name='dist-feature-test',
+    )
+    rpc_worker_names: Dict[DistRole, List[str]] = {}
+
+    rpc.init_rpc(
+        current_ctx=current_ctx,
+        rpc_worker_names=rpc_worker_names,
+        master_addr='localhost',
+        master_port=master_port,
+    )
+
+    # 2) Collect all workers:
+    partition_to_workers = rpc.rpc_partition_to_workers(
+        current_ctx, world_size, rank)
+
+    assert partition_to_workers == [
+        ['dist-feature-test-0'],
+        ['dist-feature-test-1'],
+    ]
+
+    # 3) Find the mapping between worker and partition ID:
+    rpc_router = RpcRouter(partition_to_workers)
+
+    assert rpc_router.get_to_worker(partition_idx=0) == 'dist-feature-test-0'
+    assert rpc_router.get_to_worker(partition_idx=1) == 'dist-feature-test-1'
+
+    meta = {
+        'edge_types': None,
+        'is_hetero': False,
+        'node_types': None,
+        'num_parts': 2
+    }
+
+    feature.num_partitions = world_size
+    feature.partition_idx = rank
+    feature.feature_pb = partition_book
+    feature.meta = meta
+    feature.set_local_only(local_only=False)
+    feature.set_rpc_router(rpc_router)
+
+    # Global node IDs:
+    global_id0 = torch.arange(128 * 2)
+    global_id1 = torch.arange(128 * 2) + 128 * 2
+
+    # Lookup the features from stores including locally and remotely:
+    tensor0 = feature.lookup_features(global_id0)
+    tensor1 = feature.lookup_features(global_id1)
+
+    # Expected searched results:
+    cpu_tensor0 = torch.cat([torch.ones(128, 1024), torch.ones(128, 1024) * 2])
+    cpu_tensor1 = torch.cat([torch.zeros(128, 1024), torch.zeros(128, 1024)])
+
+    # Verify..
+    assert torch.allclose(cpu_tensor0, tensor0.wait())
+    assert torch.allclose(cpu_tensor1, tensor1.wait())
+
+    rpc.shutdown_rpc()
+
+
+def test_dist_feature_lookup():
+    cpu_tensor0 = torch.cat([torch.ones(128, 1024), torch.ones(128, 1024) * 2])
+    cpu_tensor1 = torch.cat([torch.zeros(128, 1024), torch.zeros(128, 1024)])
+
+    # Global node IDs:
+    global_id0 = torch.arange(128 * 2)
+    global_id1 = torch.arange(128 * 2) + 128 * 2
+
+    # Set the partition book for two features (partition 0 and 1):
+    partition_book = torch.cat([
+        torch.zeros(128 * 2, dtype=torch.long),
+        torch.ones(128 * 2, dtype=torch.long)
+    ])
+
+    # Put the test tensor into the different feature stores with IDs:
+    feature0 = LocalFeatureStore()
+    feature0.put_global_id(global_id0, group_name=None)
+    feature0.put_tensor(cpu_tensor0, group_name=None, attr_name='x')
+
+    feature1 = LocalFeatureStore()
+    feature1.put_global_id(global_id1, group_name=None)
+    feature1.put_tensor(cpu_tensor1, group_name=None, attr_name='x')
+
+    mp_context = torch.multiprocessing.get_context('spawn')
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    s.bind(('127.0.0.1', 0))
+    port = s.getsockname()[1]
+    s.close()
+
+    w0 = mp_context.Process(target=run_rpc_feature_test,
+                            args=(2, 0, feature0, partition_book, port))
+    w1 = mp_context.Process(target=run_rpc_feature_test,
+                            args=(2, 1, feature1, partition_book, port))
+
+    w0.start()
+    w1.start()
+    w0.join()
+    w1.join()