1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
|
# Owner(s): ["oncall: distributed"]
from unittest import mock
import torch.distributed as c10d
from torch.distributed.collective_utils import all_gather, broadcast
from torch.testing._internal.common_distributed import MultiProcessTestCase
class TestCollectiveUtils(MultiProcessTestCase):
def setUp(self):
super().setUp()
self._spawn_processes()
def tearDown(self) -> None:
super().tearDown()
def opts(self, threads=2):
opts = c10d.ProcessGroupGloo._Options()
opts._timeout = 50.0
opts._threads = threads
return opts
def test_broadcast_result(self) -> None:
"""
Basic unit test for broadcast using a process group of default world size.
"""
store = c10d.FileStore(self.file_name, self.world_size)
c10d.init_process_group(
backend="gloo", store=store, rank=self.rank, world_size=self.world_size
)
pg = c10d.new_group(pg_options=self.opts())
func = mock.MagicMock()
func.return_value = pg.rank()
res = broadcast(data_or_fn=func, rank=0, pg=pg)
assert res == 0, f"Expect res to be 0 (got {res})"
if pg.rank() == 0:
func.assert_called_once()
else:
func.assert_not_called()
func.reset_mock()
res = broadcast(data_or_fn=func, rank=1, pg=pg)
assert res == 1, f"Expect res to be 1 (got {res})"
if pg.rank() == 1:
func.assert_called_once()
else:
func.assert_not_called()
def test_broadcast_result_no_pg(self) -> None:
"""
Ensure broadcast has no dependency on torch.distributed when run in single process.
"""
func = mock.MagicMock()
res = broadcast(data_or_fn=func, rank=0)
func.assert_called_once()
def test_broadcast_result_raises_exceptions_from_func(
self,
) -> None:
"""
Ensure broadcast exception is propagated properly.
"""
# no process group
func = mock.MagicMock()
exc = Exception("test exception")
func.side_effect = exc
expected_exception = "test exception"
with self.assertRaisesRegex(Exception, expected_exception):
broadcast(data_or_fn=func, rank=0)
def test_all_gather_result(self) -> None:
"""
Basic unit test for all_gather using a process group of default world size.
"""
store = c10d.FileStore(self.file_name, self.world_size)
c10d.init_process_group(
backend="gloo", store=store, rank=self.rank, world_size=self.world_size
)
pg = c10d.new_group(pg_options=self.opts())
func = mock.MagicMock()
func.return_value = pg.rank()
res = all_gather(data_or_fn=func, pg=pg)
func.assert_called_once()
assert res == list(
range(self.world_size)
), f"Expect res to be list of 0 through {self.world_size} (got {res})"
def test_all_gather_result_no_pg(self) -> None:
"""
Ensure all_gather has no dependency on torch.distributed when run in single process.
"""
func = mock.MagicMock()
res = all_gather(data_or_fn=func)
func.assert_called_once()
def test_all_gather_result_raises_exceptions_from_func(
self,
) -> None:
"""
Ensure all_gather exception is propagated properly.
"""
# no process group
func = mock.MagicMock()
exc = Exception("test exception")
func.side_effect = exc
expected_exception = "test exception"
with self.assertRaisesRegex(Exception, expected_exception):
all_gather(data_or_fn=func)
|
