Repository: pyg-team/pytorch_geometric
Pull request #8978 (2024-02-27 16:55:36)

Improved distributed tests
This PR fixes distributed tests in cases when an Exception raised in subprocess is not correctly manifested to pytest, finally reported as PASSED, causing false positive result.
Changes made: 

- `assert_run_mproc` test utility function for decorating and running Processes
- `MPCaptOutput` context utility class, capturing std output and stderr 

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8978?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: Patch coverage is `29.31034%` with `41 lines` in your changes are missing coverage. Please review.
> Project coverage is 89.15%. Comparing base [(`9b660ac`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/9b660ac6ca882604d1ae521912d20ded1d180ecf?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) to head [(`3ca6a06`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8978?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).

> :exclamation: Current head 3ca6a06 differs from pull request most recent head 2463cd6. Consider uploading reports for the commit 2463cd6 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8978?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/testing/dist\_utils.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8978?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3Rlc3RpbmcvZGlzdF91dGlscy5weQ==) | 28.07% | [41 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8978?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #8978      +/-   ##
==========================================
- Coverage   89.89%   89.15%   -0.74%     
==========================================
  Files         468      469       +1     
  Lines       30025    30078      +53     
==========================================
- Hits        26991    26817     -174     
- Misses       3034     3261     +227     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8978?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).

JakubPietrakIntel: `Notice:  A new release of pip is available: 23.0.1 -> 24.0
Notice:  To update, run: pip install --upgrade pip
torch_geometric/testing/dist_utils.py:25: error: "TextIO" has no attribute "getvalue"  [attr-defined]
torch_geometric/testing/dist_utils.py:28: error: "TextIO" has no attribute "getvalue"  [attr-defined]
torch_geometric/testing/dist_utils.py:64: error: "Process" has no attribute "_target"  [attr-defined]
torch_geometric/testing/dist_utils.py:65: error: "Process" has no attribute "_target"  [attr-defined]
torch_geometric/testing/dist_utils.py:66: error: "Process" has no attribute "_args"  [attr-defined]`

@arlesniak we need to look into thos CI issues first

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 56d72ad53..1ede8d9a4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for `EdgeIndex` in `spmm` ([#9026](https://github.com/pyg-team/pytorch_geometric/pull/9026))
 - Added option to pre-allocate memory in GPU-based `ApproxKNN` ([#9046](https://github.com/pyg-team/pytorch_geometric/pull/9046))
 - Added support for `EdgeIndex` in `MessagePassing` ([#9007](https://github.com/pyg-team/pytorch_geometric/pull/9007))
 - Added support for `torch.compile` in combination with `EdgeIndex` ([#9007](https://github.com/pyg-team/pytorch_geometric/pull/9007))
@@ -16,6 +17,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Changed distributed sampler and loader tests to correctly report failures in subprocesses to `pytest` ([#8978](https://github.com/pyg-team/pytorch_geometric/pull/8978))
 - Remove filtering of node/edge types in `trim_to_layer` functionality ([#9021](https://github.com/pyg-team/pytorch_geometric/pull/9021))
 - Default to `scatter` operations in `MessagePassing` in case `torch.use_deterministic_algorithms` is not set ([#9009](https://github.com/pyg-team/pytorch_geometric/pull/9009))
 - Made `MessagePassing` interface thread-safe ([#9001](https://github.com/pyg-team/pytorch_geometric/pull/9001))
diff --git a/torch_geometric/utils/_spmm.py b/torch_geometric/utils/_spmm.py
index ec89adda8..a0bc34cd7 100644
--- a/torch_geometric/utils/_spmm.py
+++ b/torch_geometric/utils/_spmm.py
@@ -4,6 +4,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
+from torch_geometric import EdgeIndex
 from torch_geometric.typing import Adj, SparseTensor, torch_sparse
 from torch_geometric.utils import is_torch_sparse_tensor, scatter
 
@@ -16,9 +17,11 @@ def spmm(
     r"""Matrix product of sparse matrix with dense matrix.
 
     Args:
-        src (torch.Tensor or torch_sparse.SparseTensor): The input sparse
-            matrix, either a :pyg:`PyG` :class:`torch_sparse.SparseTensor` or a
-            :pytorch:`PyTorch` :class:`torch.sparse.Tensor`.
+        src (torch.Tensor or torch_sparse.SparseTensor or EdgeIndex):
+            The input sparse matrix which can be a
+            :pyg:`PyG` :class:`torch_sparse.SparseTensor`,
+            a :pytorch:`PyTorch` :class:`torch.sparse.Tensor` or
+            a :pyg:`PyG` :class:`EdgeIndex`.
         other (torch.Tensor): The input dense matrix.
         reduce (str, optional): The reduce operation to use
             (:obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`).
@@ -31,6 +34,9 @@ def spmm(
     if reduce not in ['sum', 'mean', 'min', 'max']:
         raise ValueError(f"`reduce` argument '{reduce}' not supported")
 
+    if not torch.jit.is_scripting() and isinstance(src, EdgeIndex):
+        return src.matmul(other=other, reduce=reduce)  # type: ignore
+
     if isinstance(src, SparseTensor):
         if src.nnz() == 0:
             return other.new_zeros(src.size(0), other.size(1))


Test patch:
diff --git a/test/distributed/test_dist_link_neighbor_loader.py b/test/distributed/test_dist_link_neighbor_loader.py
index bbdc23a36..eea8e7eae 100644
--- a/test/distributed/test_dist_link_neighbor_loader.py
+++ b/test/distributed/test_dist_link_neighbor_loader.py
@@ -16,6 +16,7 @@ from torch_geometric.distributed import (
     Partitioner,
 )
 from torch_geometric.testing import onlyDistributedTest, withMETIS
+from torch_geometric.testing.distributed import ProcArgs, assert_run_mproc
 
 
 def create_dist_data(tmp_path: str, rank: int):
@@ -26,8 +27,8 @@ def create_dist_data(tmp_path: str, rank: int):
 
 
 def dist_link_neighbor_loader_homo(
-    tmp_path: str,
     world_size: int,
+    tmp_path: str,
     rank: int,
     master_addr: str,
     master_port: int,
@@ -76,8 +77,8 @@ def dist_link_neighbor_loader_homo(
 
 
 def dist_link_neighbor_loader_hetero(
-    tmp_path: str,
     world_size: int,
+    tmp_path: str,
     rank: int,
     master_addr: str,
     master_port: int,
@@ -166,22 +167,14 @@ def test_dist_link_neighbor_loader_homo(
     partitioner = Partitioner(data, num_parts, tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_link_neighbor_loader_homo,
-        args=(tmp_path, num_parts, 0, addr, port, num_workers, async_sampling,
-              neg_ratio),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_link_neighbor_loader_homo,
-        args=(tmp_path, num_parts, 1, addr, port, num_workers, async_sampling,
-              neg_ratio),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(
+            target=dist_link_neighbor_loader_homo,
+            args=(tmp_path, part, addr, port, num_workers, async_sampling,
+                  neg_ratio),
+        ) for part in range(num_parts)
+    ]
+    assert_run_mproc(mp_context, procs)
 
 
 @withMETIS
@@ -218,19 +211,11 @@ def test_dist_link_neighbor_loader_hetero(
     partitioner = Partitioner(data, num_parts, tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_link_neighbor_loader_hetero,
-        args=(tmp_path, num_parts, 0, addr, port, num_workers, async_sampling,
-              neg_ratio, edge_type),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_link_neighbor_loader_hetero,
-        args=(tmp_path, num_parts, 1, addr, port, num_workers, async_sampling,
-              neg_ratio, edge_type),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(
+            target=dist_link_neighbor_loader_hetero,
+            args=(tmp_path, part, addr, port, num_workers, async_sampling,
+                  neg_ratio, edge_type),
+        ) for part in range(num_parts)
+    ]
+    assert_run_mproc(mp_context, procs)
diff --git a/test/distributed/test_dist_link_neighbor_sampler.py b/test/distributed/test_dist_link_neighbor_sampler.py
index a55634ac4..5915b7b16 100644
--- a/test/distributed/test_dist_link_neighbor_sampler.py
+++ b/test/distributed/test_dist_link_neighbor_sampler.py
@@ -20,6 +20,7 @@ from torch_geometric.distributed.rpc import init_rpc, shutdown_rpc
 from torch_geometric.sampler import EdgeSamplerInput, NeighborSampler
 from torch_geometric.sampler.neighbor_sampler import edge_sample
 from torch_geometric.testing import onlyDistributedTest, withMETIS
+from torch_geometric.testing.distributed import ProcArgs, assert_run_mproc
 from torch_geometric.typing import EdgeType
 
 
@@ -249,9 +250,9 @@ def dist_link_neighbor_sampler_temporal(
 
 
 def dist_link_neighbor_sampler_hetero(
+    world_size: int,
     data: FakeHeteroDataset,
     tmp_path: str,
-    world_size: int,
     rank: int,
     master_port: int,
     input_type: EdgeType,
@@ -340,9 +341,9 @@ def dist_link_neighbor_sampler_hetero(
 
 
 def dist_link_neighbor_sampler_temporal_hetero(
+    world_size: int,
     data: FakeHeteroDataset,
     tmp_path: str,
-    world_size: int,
     rank: int,
     master_port: int,
     input_type: EdgeType,
@@ -445,21 +446,11 @@ def test_dist_link_neighbor_sampler(disjoint):
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
-    w0 = mp_context.Process(
-        target=dist_link_neighbor_sampler,
-        args=(world_size, 0, port, disjoint),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_link_neighbor_sampler,
-        args=(world_size, 1, port, disjoint),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(target=dist_link_neighbor_sampler, args=(0, port, disjoint)),
+        ProcArgs(target=dist_link_neighbor_sampler, args=(1, port, disjoint)),
+    ]
+    assert_run_mproc(mp_context, procs)
 
 
 @onlyDistributedTest
@@ -473,21 +464,17 @@ def test_dist_link_neighbor_sampler_temporal(seed_time, temporal_strategy):
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
-    w0 = mp_context.Process(
-        target=dist_link_neighbor_sampler_temporal,
-        args=(world_size, 0, port, seed_time, temporal_strategy, 'time'),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_link_neighbor_sampler_temporal,
-        args=(world_size, 1, port, seed_time, temporal_strategy, 'time'),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(
+            target=dist_link_neighbor_sampler_temporal,
+            args=(0, port, seed_time, temporal_strategy, 'time'),
+        ),
+        ProcArgs(
+            target=dist_link_neighbor_sampler_temporal,
+            args=(1, port, seed_time, temporal_strategy, 'time'),
+        ),
+    ]
+    assert_run_mproc(mp_context, procs)
 
 
 @onlyDistributedTest
@@ -506,21 +493,17 @@ def test_dist_link_neighbor_sampler_edge_level_temporal(
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
-    w0 = mp_context.Process(
-        target=dist_link_neighbor_sampler_temporal,
-        args=(world_size, 0, port, seed_time, temporal_strategy, 'edge_time'),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_link_neighbor_sampler_temporal,
-        args=(world_size, 1, port, seed_time, temporal_strategy, 'edge_time'),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(
+            target=dist_link_neighbor_sampler_temporal,
+            args=(0, port, seed_time, temporal_strategy, 'edge_time'),
+        ),
+        ProcArgs(
+            target=dist_link_neighbor_sampler_temporal,
+            args=(1, port, seed_time, temporal_strategy, 'edge_time'),
+        ),
+    ]
+    assert_run_mproc(mp_context, procs)
 
 
 @withMETIS
@@ -534,7 +517,6 @@ def test_dist_link_neighbor_sampler_hetero(tmp_path, disjoint):
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
     data = FakeHeteroDataset(
         num_graphs=1,
         avg_num_nodes=100,
@@ -545,25 +527,21 @@ def test_dist_link_neighbor_sampler_hetero(tmp_path, disjoint):
     )[0]
     data = T.ToUndirected()(data)
 
-    partitioner = Partitioner(data, world_size, tmp_path)
+    procs = [
+        ProcArgs(
+            target=dist_link_neighbor_sampler_hetero,
+            args=(data, tmp_path, 0, port, ('v0', 'e0', 'v0'), disjoint),
+        ),
+        ProcArgs(
+            target=dist_link_neighbor_sampler_hetero,
+            args=(data, tmp_path, 1, port, ('v1', 'e0', 'v0'), disjoint),
+        ),
+    ]
+
+    partitioner = Partitioner(data, len(procs), tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_link_neighbor_sampler_hetero,
-        args=(data, tmp_path, world_size, 0, port, ('v0', 'e0', 'v0'),
-              disjoint),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_link_neighbor_sampler_hetero,
-        args=(data, tmp_path, world_size, 1, port, ('v1', 'e0', 'v0'),
-              disjoint),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    assert_run_mproc(mp_context, procs)
 
 
 @withMETIS
@@ -585,7 +563,6 @@ def test_dist_link_neighbor_sampler_temporal_hetero(
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
     data = FakeHeteroDataset(
         num_graphs=1,
         avg_num_nodes=100,
@@ -600,25 +577,23 @@ def test_dist_link_neighbor_sampler_temporal_hetero(
     data['v0'].time = torch.ones(data['v0'].num_nodes, dtype=torch.int64)
     data['v1'].time = torch.full((data['v1'].num_nodes, ), 2).long()
 
-    partitioner = Partitioner(data, world_size, tmp_path)
+    procs = [
+        ProcArgs(
+            target=dist_link_neighbor_sampler_temporal_hetero,
+            args=(data, tmp_path, 0, port, ('v0', 'e0', 'v0'), seed_time,
+                  temporal_strategy, 'time'),
+        ),
+        ProcArgs(
+            target=dist_link_neighbor_sampler_temporal_hetero,
+            args=(data, tmp_path, 1, port, ('v1', 'e0', 'v0'), seed_time,
+                  temporal_strategy, 'time'),
+        ),
+    ]
+
+    partitioner = Partitioner(data, len(procs), tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_link_neighbor_sampler_temporal_hetero,
-        args=(data, tmp_path, world_size, 0, port, ('v0', 'e0', 'v0'),
-              seed_time, temporal_strategy, 'time'),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_link_neighbor_sampler_temporal_hetero,
-        args=(data, tmp_path, world_size, 1, port, ('v1', 'e0', 'v0'),
-              seed_time, temporal_strategy, 'time'),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    assert_run_mproc(mp_context, procs)
 
 
 @withMETIS
@@ -639,7 +614,6 @@ def test_dist_link_neighbor_sampler_edge_level_temporal_hetero(
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
     data = FakeHeteroDataset(
         num_graphs=1,
         avg_num_nodes=100,
@@ -655,22 +629,20 @@ def test_dist_link_neighbor_sampler_edge_level_temporal_hetero(
         data[edge_type].edge_time = torch.full(  #
             (data[edge_type].num_edges, ), i, dtype=torch.int64)
 
-    partitioner = Partitioner(data, world_size, tmp_path)
+    procs = [
+        ProcArgs(
+            target=dist_link_neighbor_sampler_temporal_hetero,
+            args=(data, tmp_path, 0, port, ('v0', 'e0', 'v0'), seed_time,
+                  temporal_strategy, 'edge_time'),
+        ),
+        ProcArgs(
+            target=dist_link_neighbor_sampler_temporal_hetero,
+            args=(data, tmp_path, 1, port, ('v0', 'e0', 'v1'), seed_time,
+                  temporal_strategy, 'edge_time'),
+        ),
+    ]
+
+    partitioner = Partitioner(data, len(procs), tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_link_neighbor_sampler_temporal_hetero,
-        args=(data, tmp_path, world_size, 0, port, ('v0', 'e0', 'v0'),
-              seed_time, temporal_strategy, 'edge_time'),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_link_neighbor_sampler_temporal_hetero,
-        args=(data, tmp_path, world_size, 1, port, ('v0', 'e0', 'v1'),
-              seed_time, temporal_strategy, 'edge_time'),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    assert_run_mproc(mp_context, procs)
diff --git a/test/distributed/test_dist_neighbor_loader.py b/test/distributed/test_dist_neighbor_loader.py
index a678cd293..c07f8cff7 100644
--- a/test/distributed/test_dist_neighbor_loader.py
+++ b/test/distributed/test_dist_neighbor_loader.py
@@ -15,6 +15,7 @@ from torch_geometric.distributed import (
     Partitioner,
 )
 from torch_geometric.testing import onlyDistributedTest, withMETIS
+from torch_geometric.testing.distributed import ProcArgs, assert_run_mproc
 
 
 def create_dist_data(tmp_path: str, rank: int):
@@ -25,8 +26,8 @@ def create_dist_data(tmp_path: str, rank: int):
 
 
 def dist_neighbor_loader_homo(
-    tmp_path: str,
     world_size: int,
+    tmp_path: str,
     rank: int,
     master_addr: str,
     master_port: int,
@@ -78,8 +79,8 @@ def dist_neighbor_loader_homo(
 
 
 def dist_neighbor_loader_hetero(
-    tmp_path: str,
     world_size: int,
+    tmp_path: str,
     rank: int,
     master_addr: str,
     master_port: int,
@@ -170,20 +171,13 @@ def test_dist_neighbor_loader_homo(
     partitioner = Partitioner(data, num_parts, tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_neighbor_loader_homo,
-        args=(tmp_path, num_parts, 0, addr, port, num_workers, async_sampling),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_neighbor_loader_homo,
-        args=(tmp_path, num_parts, 1, addr, port, num_workers, async_sampling),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(
+            target=dist_neighbor_loader_homo,
+            args=(tmp_path, part, addr, port, num_workers, async_sampling),
+        ) for part in range(num_parts)
+    ]
+    assert_run_mproc(mp_context, procs)
 
 
 @withMETIS
@@ -216,17 +210,10 @@ def test_dist_neighbor_loader_hetero(
     partitioner = Partitioner(data, num_parts, tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_neighbor_loader_hetero,
-        args=(tmp_path, num_parts, 0, addr, port, num_workers, async_sampling),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_neighbor_loader_hetero,
-        args=(tmp_path, num_parts, 1, addr, port, num_workers, async_sampling),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(
+            target=dist_neighbor_loader_hetero,
+            args=(tmp_path, part, addr, port, num_workers, async_sampling),
+        ) for part in range(num_parts)
+    ]
+    assert_run_mproc(mp_context, procs)
diff --git a/test/distributed/test_dist_neighbor_sampler.py b/test/distributed/test_dist_neighbor_sampler.py
index fbea1f23e..7ed2e035f 100644
--- a/test/distributed/test_dist_neighbor_sampler.py
+++ b/test/distributed/test_dist_neighbor_sampler.py
@@ -19,6 +19,7 @@ from torch_geometric.distributed.rpc import init_rpc, shutdown_rpc
 from torch_geometric.sampler import NeighborSampler, NodeSamplerInput
 from torch_geometric.sampler.neighbor_sampler import node_sample
 from torch_geometric.testing import onlyDistributedTest, withMETIS
+from torch_geometric.testing.distributed import ProcArgs, assert_run_mproc
 
 
 def create_data(rank: int, world_size: int, time_attr: Optional[str] = None):
@@ -225,9 +226,9 @@ def dist_neighbor_sampler_temporal(
 
 
 def dist_neighbor_sampler_hetero(
+    world_size: int,
     data: FakeHeteroDataset,
     tmp_path: str,
-    world_size: int,
     rank: int,
     master_port: int,
     input_type: str,
@@ -304,9 +305,9 @@ def dist_neighbor_sampler_hetero(
 
 
 def dist_neighbor_sampler_temporal_hetero(
+    world_size: int,
     data: FakeHeteroDataset,
     tmp_path: str,
-    world_size: int,
     rank: int,
     master_port: int,
     input_type: str,
@@ -396,21 +397,11 @@ def test_dist_neighbor_sampler(disjoint):
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
-    w0 = mp_context.Process(
-        target=dist_neighbor_sampler,
-        args=(world_size, 0, port, disjoint),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_neighbor_sampler,
-        args=(world_size, 1, port, disjoint),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(target=dist_neighbor_sampler, args=(0, port, disjoint)),
+        ProcArgs(target=dist_neighbor_sampler, args=(1, port, disjoint)),
+    ]
+    assert_run_mproc(mp_context, procs)
 
 
 @onlyDistributedTest
@@ -424,21 +415,17 @@ def test_dist_neighbor_sampler_temporal(seed_time, temporal_strategy):
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
-    w0 = mp_context.Process(
-        target=dist_neighbor_sampler_temporal,
-        args=(world_size, 0, port, seed_time, temporal_strategy, 'time'),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_neighbor_sampler_temporal,
-        args=(world_size, 1, port, seed_time, temporal_strategy, 'time'),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(
+            target=dist_neighbor_sampler_temporal,
+            args=(0, port, seed_time, temporal_strategy, 'time'),
+        ),
+        ProcArgs(
+            target=dist_neighbor_sampler_temporal,
+            args=(1, port, seed_time, temporal_strategy, 'time'),
+        ),
+    ]
+    assert_run_mproc(mp_context, procs)
 
 
 @onlyDistributedTest
@@ -457,21 +444,17 @@ def test_dist_neighbor_sampler_edge_level_temporal(
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
-    w0 = mp_context.Process(
-        target=dist_neighbor_sampler_temporal,
-        args=(world_size, 0, port, seed_time, temporal_strategy, 'edge_time'),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_neighbor_sampler_temporal,
-        args=(world_size, 1, port, seed_time, temporal_strategy, 'edge_time'),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    procs = [
+        ProcArgs(
+            target=dist_neighbor_sampler_temporal,
+            args=(0, port, seed_time, temporal_strategy, 'edge_time'),
+        ),
+        ProcArgs(
+            target=dist_neighbor_sampler_temporal,
+            args=(1, port, seed_time, temporal_strategy, 'edge_time'),
+        ),
+    ]
+    assert_run_mproc(mp_context, procs)
 
 
 @withMETIS
@@ -485,7 +468,6 @@ def test_dist_neighbor_sampler_hetero(tmp_path, disjoint):
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
     data = FakeHeteroDataset(
         num_graphs=1,
         avg_num_nodes=100,
@@ -495,23 +477,21 @@ def test_dist_neighbor_sampler_hetero(tmp_path, disjoint):
         edge_dim=2,
     )[0]
 
-    partitioner = Partitioner(data, world_size, tmp_path)
+    procs = [
+        ProcArgs(
+            target=dist_neighbor_sampler_hetero,
+            args=(data, tmp_path, 0, port, 'v0', disjoint),
+        ),
+        ProcArgs(
+            target=dist_neighbor_sampler_hetero,
+            args=(data, tmp_path, 1, port, 'v1', disjoint),
+        ),
+    ]
+
+    partitioner = Partitioner(data, len(procs), tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_neighbor_sampler_hetero,
-        args=(data, tmp_path, world_size, 0, port, 'v0', disjoint),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_neighbor_sampler_hetero,
-        args=(data, tmp_path, world_size, 1, port, 'v1', disjoint),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    assert_run_mproc(mp_context, procs)
 
 
 @withMETIS
@@ -533,7 +513,6 @@ def test_dist_neighbor_sampler_temporal_hetero(
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
     data = FakeHeteroDataset(
         num_graphs=1,
         avg_num_nodes=100,
@@ -548,25 +527,23 @@ def test_dist_neighbor_sampler_temporal_hetero(
     data['v1'].time = torch.full((data.num_nodes_dict['v1'], ), 2,
                                  dtype=torch.int64)
 
-    partitioner = Partitioner(data, world_size, tmp_path)
+    procs = [
+        ProcArgs(
+            target=dist_neighbor_sampler_temporal_hetero,
+            args=(data, tmp_path, 0, port, 'v0', seed_time, temporal_strategy,
+                  'time'),
+        ),
+        ProcArgs(
+            target=dist_neighbor_sampler_temporal_hetero,
+            args=(data, tmp_path, 1, port, 'v1', seed_time, temporal_strategy,
+                  'time'),
+        ),
+    ]
+
+    partitioner = Partitioner(data, len(procs), tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_neighbor_sampler_temporal_hetero,
-        args=(data, tmp_path, world_size, 0, port, 'v0', seed_time,
-              temporal_strategy, 'time'),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_neighbor_sampler_temporal_hetero,
-        args=(data, tmp_path, world_size, 1, port, 'v1', seed_time,
-              temporal_strategy, 'time'),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    assert_run_mproc(mp_context, procs)
 
 
 @withMETIS
@@ -587,7 +564,6 @@ def test_dist_neighbor_sampler_edge_level_temporal_hetero(
         s.bind(('', 0))
         port = s.getsockname()[1]
 
-    world_size = 2
     data = FakeHeteroDataset(
         num_graphs=1,
         avg_num_nodes=100,
@@ -601,22 +577,20 @@ def test_dist_neighbor_sampler_edge_level_temporal_hetero(
         data[edge_type].edge_time = torch.full(
             (data[edge_type].edge_index.size(1), ), i, dtype=torch.int64)
 
-    partitioner = Partitioner(data, world_size, tmp_path)
+    procs = [
+        ProcArgs(
+            target=dist_neighbor_sampler_temporal_hetero,
+            args=(data, tmp_path, 0, port, 'v0', seed_time, temporal_strategy,
+                  'edge_time'),
+        ),
+        ProcArgs(
+            target=dist_neighbor_sampler_temporal_hetero,
+            args=(data, tmp_path, 1, port, 'v1', seed_time, temporal_strategy,
+                  'edge_time'),
+        ),
+    ]
+
+    partitioner = Partitioner(data, len(procs), tmp_path)
     partitioner.generate_partition()
 
-    w0 = mp_context.Process(
-        target=dist_neighbor_sampler_temporal_hetero,
-        args=(data, tmp_path, world_size, 0, port, 'v0', seed_time,
-              temporal_strategy, 'edge_time'),
-    )
-
-    w1 = mp_context.Process(
-        target=dist_neighbor_sampler_temporal_hetero,
-        args=(data, tmp_path, world_size, 1, port, 'v1', seed_time,
-              temporal_strategy, 'edge_time'),
-    )
-
-    w0.start()
-    w1.start()
-    w0.join()
-    w1.join()
+    assert_run_mproc(mp_context, procs)
diff --git a/test/utils/test_spmm.py b/test/utils/test_spmm.py
index a856a3cb3..76cc8cd94 100644
--- a/test/utils/test_spmm.py
+++ b/test/utils/test_spmm.py
@@ -6,6 +6,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
+from torch_geometric import EdgeIndex
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import withCUDA, withPackage
 from torch_geometric.typing import SparseTensor
@@ -106,6 +107,24 @@ def test_spmm_jit(reduce):
         assert torch.allclose(out2, out3, atol=1e-6)
 
 
+@withCUDA
+@withPackage('torch>=2.0.0')
+@pytest.mark.parametrize('reduce', ['sum', 'mean', 'min', 'max'])
+def test_spmm_edge_index(device, reduce):
+    src = EdgeIndex(
+        [[0, 1, 1, 2], [1, 0, 2, 1]],
+        sparse_size=(4, 3),
+        sort_order='row',
+        device=device,
+    )
+    other = torch.rand(3, 4, device=device)
+    out = spmm(src, other, reduce=reduce)
+    assert out.size() == (4, 4)
+
+    out2 = spmm(src.to_sparse_coo(), other, reduce=reduce)
+    assert torch.allclose(out, out2)
+
+
 if __name__ == '__main__':
     import argparse
 
diff --git a/torch_geometric/testing/distributed.py b/torch_geometric/testing/distributed.py
new file mode 100644
index 000000000..d96699a4c
--- /dev/null
+++ b/torch_geometric/testing/distributed.py
@@ -0,0 +1,92 @@
+import sys
+import traceback
+from dataclasses import dataclass
+from io import StringIO
+from typing import Any, Callable, List, Tuple
+
+import pytest
+from torch.multiprocessing import Manager, Queue
+from typing_extensions import Self
+
+
+@dataclass
+class ProcArgs:
+    target: Callable
+    args: Tuple[Any, ...]
+
+
+class MPCaptOutput:
+    def __enter__(self) -> Self:
+        self.stdout = StringIO()
+        self.stderr = StringIO()
+        self.old_stdout = sys.stdout
+        self.old_stderr = sys.stderr
+
+        sys.stdout = self.stdout
+        sys.stderr = self.stderr
+
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        sys.stdout = self.old_stdout
+        sys.stderr = self.old_stderr
+
+    @property
+    def stdout_str(self) -> str:
+        return self.stdout.getvalue()
+
+    @property
+    def stderr_str(self) -> str:
+        return self.stderr.getvalue()
+
+
+def ps_std_capture(
+    func: Callable,
+    queue: Queue,
+    *args: Any,
+    **kwargs: Any,
+) -> None:
+    with MPCaptOutput() as capt:
+        try:
+            func(*args, **kwargs)
+        except Exception as e:
+            traceback.print_exc(file=sys.stderr)
+            raise e
+        finally:
+            queue.put((capt.stdout_str, capt.stderr_str))
+
+
+def assert_run_mproc(
+    mp_context: Any,
+    pargs: List[ProcArgs],
+    full_trace: bool = False,
+    timeout: int = 5,
+) -> None:
+    manager = Manager()
+    world_size = len(pargs)
+    queues = [manager.Queue() for _ in pargs]
+    procs = [
+        mp_context.Process(
+            target=ps_std_capture,
+            args=[p.target, q, world_size] + list(p.args),
+        ) for p, q in zip(pargs, queues)
+    ]
+    results = []
+
+    for p, q in zip(procs, queues):
+        p.start()
+
+    for p, q in zip(procs, queues):
+        p.join()
+        stdout, stderr = q.get(timeout=timeout)
+        results.append((p, stdout, stderr))
+
+    for p, stdout, stderr in results:
+        if stdout:
+            print(stdout)
+        if stderr:  # can be a warning as well => exitcode == 0
+            print(stderr)
+        if p.exitcode != 0:
+            pytest.fail(
+                pytrace=full_trace, reason=stderr.splitlines()[-1]
+                if stderr else f"exitcode {p.exitcode}")