-
Notifications
You must be signed in to change notification settings - Fork 21.4k
/
common_distributed.py
455 lines (394 loc) · 16.5 KB
/
common_distributed.py
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
from multiprocessing import Manager
from contextlib import contextmanager
from io import StringIO
import os
import sys
import tempfile
import time
import unittest
import logging
import traceback
import types
from typing import NamedTuple
from functools import wraps
import torch
import torch.distributed as c10d
from functools import partial, reduce
from torch.testing._internal.common_utils import TestCase, TEST_WITH_ROCM, FILE_SCHEMA
class TestSkip(NamedTuple):
exit_code: int
message: str
TEST_SKIPS = {
"backend_unavailable": TestSkip(72, "Skipped because distributed backend is not available."),
"small_worldsize": TestSkip(73, "Skipped due to small world size."),
"no_cuda": TestSkip(74, "CUDA is not available."),
"multi-gpu": TestSkip(75, "Need at least 2 CUDA devices"),
"nccl": TestSkip(76, "c10d not compiled with NCCL support"),
"skipIfRocm": TestSkip(78, "Test skipped for ROCm")
}
def skip_if_no_gpu(func):
""" Nccl multigpu tests require at least 2 GPUS. Skip if this is not met"""
@wraps(func)
def wrapper(*args, **kwargs):
if not torch.cuda.is_available():
sys.exit(TEST_SKIPS["no_cuda"].exit_code)
if torch.cuda.device_count() < int(os.environ["WORLD_SIZE"]):
message = "Need at least {} CUDA devices".format(os.environ["WORLD_SIZE"])
TEST_SKIPS["multi-gpu"] = TestSkip(75, message)
sys.exit(TEST_SKIPS["multi-gpu"].exit_code)
return func(*args, **kwargs)
return wrapper
def skip_if_small_worldsize(func):
@wraps(func)
def wrapper(*args, **kwargs):
if (os.environ["BACKEND"] != "mpi") and int(os.environ["WORLD_SIZE"]) <= 2:
sys.exit(TEST_SKIPS["small_worldsize"].exit_code)
return func(*args, **kwargs)
return wrapper
def skip_if_not_multigpu(func):
"""Multi-GPU tests requires at least 2 GPUS. Skip if this is not met."""
@wraps(func)
def wrapper(*args, **kwargs):
if torch.cuda.is_available() and torch.cuda.device_count() >= 2:
return func(*args, **kwargs)
message = "Need at least {} CUDA devices".format(2)
TEST_SKIPS["multi-gpu"] = TestSkip(75, message)
sys.exit(TEST_SKIPS['multi-gpu'].exit_code)
return wrapper
def require_n_gpus_for_nccl_backend(n, backend):
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
if backend == "nccl" and torch.cuda.device_count() < n:
message = "Need at least {} CUDA devices".format(n)
TEST_SKIPS["multi-gpu"] = TestSkip(75, message)
sys.exit(TEST_SKIPS['multi-gpu'].exit_code)
else:
return func(*args, **kwargs)
return wrapper
return decorator
def skip_if_lt_x_gpu(x):
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
if torch.cuda.is_available() and torch.cuda.device_count() >= x:
return func(*args, **kwargs)
message = "Need at least {} CUDA devices".format(x)
TEST_SKIPS["multi-gpu"] = TestSkip(75, message)
sys.exit(TEST_SKIPS['multi-gpu'].exit_code)
return wrapper
return decorator
def requires_gloo():
return unittest.skipUnless(
c10d.is_gloo_available(),
"c10d was not compiled with the Gloo backend",
)
def requires_nccl_version(version, msg):
if not c10d.is_nccl_available():
return unittest.skip(
"c10d was not compiled with the NCCL backend",
)
else:
return unittest.skipIf(
torch.cuda.nccl.version() < version,
"Requires NCCL version greater than or equal to: {}, found: {}, reason: {}".format(
version,
torch.cuda.nccl.version(), msg),
)
def requires_nccl():
return unittest.skipUnless(
c10d.is_nccl_available(),
"c10d was not compiled with the NCCL backend",
)
def requires_mpi():
return unittest.skipUnless(
c10d.is_mpi_available(),
"c10d was not compiled with the MPI backend",
)
def skip_if_rocm_single_process(func):
"""Skips a test for ROCm in a single process environment"""
func.skip_if_rocm = True
@wraps(func)
def wrapper(*args, **kwargs):
if not TEST_WITH_ROCM:
return func(*args, **kwargs)
raise unittest.SkipTest("Test skipped for ROCm")
return wrapper
def skip_if_rocm(func):
"""Skips a test for ROCm"""
func.skip_if_rocm = True
@wraps(func)
def wrapper(*args, **kwargs):
if not TEST_WITH_ROCM:
return func(*args, **kwargs)
sys.exit(TEST_SKIPS['skipIfRocm'].exit_code)
return wrapper
def skip_if_win32():
return unittest.skipIf(
sys.platform == 'win32',
"This unit test case is not supportted on Windows platform",
)
TIMEOUT_DEFAULT = 100
TIMEOUT_OVERRIDE = {"test_ddp_uneven_inputs": 400}
def create_device(interface=None):
if sys.platform == 'win32' or interface is None:
return c10d.ProcessGroupGloo.create_device(hostname="127.0.0.1")
else:
return c10d.ProcessGroupGloo.create_device(interface=interface)
def get_timeout(test_id):
return TIMEOUT_OVERRIDE.get(test_id.split('.')[-1], TIMEOUT_DEFAULT)
@contextmanager
def captured_output():
new_out, new_err = StringIO(), StringIO()
old_out, old_err = sys.stdout, sys.stderr
try:
sys.stdout, sys.stderr = new_out, new_err
yield sys.stdout, sys.stderr
finally:
sys.stdout, sys.stderr = old_out, old_err
def simple_sparse_reduce_tests(rank, world_size, num_inputs=1):
"""
Generate a number of basic test cases for sparse reduction.
These cover tensors with a varying number of sparse dimensions and a varying
number of dense dimensions. The only reduction operation we support is sum.
"""
def generate(rank, world_size, sparse_dims=1, dense_dims=0):
# First sparse dimension is [0..rank].
# Subsequent dimensions are always 0, so we know there is
# a non-empty intersection between any two sparse tensors.
indices = torch.reshape(torch.arange(rank + 1), (1, rank + 1))
shape = [world_size] + [2 for _ in range(dense_dims)]
for _ in range(sparse_dims - 1):
indices = torch.cat((indices, torch.zeros(1, rank + 1)))
shape.append(world_size)
values = torch.ones([rank + 1] + [2 for _ in range(dense_dims)])
return torch.sparse_coo_tensor(indices, values, shape)
def compute_sum(fn, world_size):
return reduce(lambda a, b: a + b, [fn(rank, world_size) for rank in range(world_size)])
return [
(
[
fn(num_inputs * rank + i, num_inputs * world_size)
for i in range(num_inputs)
],
[
compute_sum(fn, num_inputs * world_size)
for i in range(num_inputs)
],
)
for fn in [
partial(generate, sparse_dims=1),
partial(generate, sparse_dims=2),
partial(generate, sparse_dims=3),
partial(generate, dense_dims=1),
partial(generate, dense_dims=2),
partial(generate, dense_dims=3),
]
]
tmp_dir = None
def initialize_temp_directories(init_method=None):
global tmp_dir
tmp_dir = tempfile.TemporaryDirectory()
os.environ["TEMP_DIR"] = tmp_dir.name
os.mkdir(os.path.join(tmp_dir.name, "barrier"))
os.mkdir(os.path.join(tmp_dir.name, "test_dir"))
init_dir_path = os.path.join(tmp_dir.name, "init_dir")
os.mkdir(init_dir_path)
# Set init method if specified.
if init_method is not None:
os.environ["INIT_METHOD"] = init_method
else:
os.environ["INIT_METHOD"] = FILE_SCHEMA + os.path.join(
init_dir_path, "shared_init_file"
)
def cleanup_temp_dir():
if tmp_dir is not None:
tmp_dir.cleanup()
# [How does MultiProcessTestCase work?]
# Each MultiProcessTestCase instance uses 1 + `world_size()` processes, by
# default `world_size()` returns 4. Let's take `test_rpc_spawn.py` as an
# example which inherits from this class. Its `Setup()` methods calls into
# `MultiProcessTestCase._spawn_processes()` which spawns `world_size()`
# subprocesses. During the spawn, the main process passes the test name to
# subprocesses, and the name is acquired from self.id(). The subprocesses
# then use the provided test function name to retrieve the function attribute
# from the test instance and run it. The main process simply waits for all
# subprocesses to join.
class MultiProcessTestCase(TestCase):
MAIN_PROCESS_RANK = -1
# This exit code is used to indicate that the test code had an error and
# exited abnormally. There are certain tests that might use sys.exit() to
# simulate failures and in those cases, we can't have an exit code of 0,
# but we still want to ensure we didn't run into any other errors.
TEST_ERROR_EXIT_CODE = 10
@property
def world_size(self):
return 4
def join_or_run(self, fn):
@wraps(fn)
def wrapper(self):
if self.rank == self.MAIN_PROCESS_RANK:
self._join_processes(fn)
else:
try:
fn()
except Exception as e:
logging.error('Caught exception: \n{}exiting process with exit code: {}'
.format(traceback.format_exc(), MultiProcessTestCase.TEST_ERROR_EXIT_CODE))
sys.exit(MultiProcessTestCase.TEST_ERROR_EXIT_CODE)
return types.MethodType(wrapper, self)
# The main process spawns N subprocesses that run the test.
# Constructor patches current instance test method to
# assume the role of the main process and join its subprocesses,
# or run the underlying test function.
def __init__(self, method_name='runTest'):
super().__init__(method_name)
fn = getattr(self, method_name)
setattr(self, method_name, self.join_or_run(fn))
def setUp(self):
super().setUp()
self.skip_return_code_checks = []
self.processes = []
self.rank = self.MAIN_PROCESS_RANK
self.file_name = tempfile.NamedTemporaryFile(delete=False).name
global TEST_SKIPS
self.old_test_skips = TEST_SKIPS.copy()
def tearDown(self):
super().tearDown()
for p in self.processes:
p.terminate()
# Each Process instance holds a few open file descriptors. The unittest
# runner creates a new TestCase instance for each test method and keeps
# it alive until the end of the entire suite. We must thus reset the
# processes to prevent an effective file descriptor leak.
self.processes = []
def _current_test_name(self):
# self.id() == e.g. '__main__.TestDistributed.TestAdditive.test_get_rank'
return self.id().split(".")[-1]
def _start_processes(self, proc):
test_skips_manager = Manager()
test_skips = test_skips_manager.dict()
global TEST_SKIPS
test_skips.update(TEST_SKIPS)
TEST_SKIPS = test_skips
self.processes = []
for rank in range(int(self.world_size)):
process = proc(
target=self.__class__._run,
name='process ' + str(rank),
args=(rank, self._current_test_name(), self.file_name))
process.start()
self.processes.append(process)
def _fork_processes(self):
proc = torch.multiprocessing.get_context("fork").Process
self._start_processes(proc)
def _spawn_processes(self):
proc = torch.multiprocessing.get_context("spawn").Process
self._start_processes(proc)
@classmethod
def _run(cls, rank, test_name, file_name):
self = cls(test_name)
self.rank = rank
self.file_name = file_name
# self.id() == e.g. '__main__.TestDistributed.test_get_rank'
# We're retrieving a corresponding test and executing it.
getattr(self, test_name)()
# exit to avoid run teardown() for fork processes
sys.exit(0)
def _join_processes(self, fn):
timeout = get_timeout(self.id())
start_time = time.time()
subprocess_error = False
try:
while True:
# check to see if any subprocess exited with an error early.
for (i, p) in enumerate(self.processes):
# This is the exit code processes exit with if they
# encountered an exception.
if p.exitcode == MultiProcessTestCase.TEST_ERROR_EXIT_CODE:
print("Process {} terminated with exit code {}, terminating remaining processes.".format(i, p.exitcode))
active_children = torch.multiprocessing.active_children()
for ac in active_children:
ac.terminate()
subprocess_error = True
break
if subprocess_error:
break
# All processes have joined cleanly if they all a valid exitcode
if all([p.exitcode is not None for p in self.processes]):
break
# Check if we should time out the test. If so, we terminate each process.
elapsed = time.time() - start_time
if elapsed > timeout:
print(
"Timing out after {} seconds and killing subprocesses.".format(
timeout
)
)
for p in self.processes:
p.terminate()
break
# Sleep to avoid excessive busy polling.
time.sleep(0.1)
elapsed_time = time.time() - start_time
if fn in self.skip_return_code_checks:
self._check_no_test_errors(elapsed_time)
else:
self._check_return_codes(elapsed_time)
finally:
global TEST_SKIPS
TEST_SKIPS = self.old_test_skips
def _check_no_test_errors(self, elapsed_time):
"""
Checks that we didn't have any errors thrown in the child processes.
"""
for i, p in enumerate(self.processes):
if p.exitcode is None:
raise RuntimeError('Process {} timed out after {} seconds'.format(i, elapsed_time))
self.assertNotEqual(self.TEST_ERROR_EXIT_CODE, p.exitcode)
def _check_return_codes(self, elapsed_time):
"""
Checks that the return codes of all spawned processes match, and skips
tests if they returned a return code indicating a skipping condition.
"""
first_process = self.processes[0]
# first, we check if there are errors in actual processes
# (via TEST_ERROR_EXIT CODE), and raise an exception for those.
# the reason we do this is to attempt to raise a more helpful error
# message than "Process x terminated/timed out"
# TODO: we should pipe the exception of the failed subprocess here.
# Currently, the actual exception is displayed as a logging output.
errored_processes = [
(i, p)
for i, p in enumerate(self.processes)
if p.exitcode == MultiProcessTestCase.TEST_ERROR_EXIT_CODE
]
if errored_processes:
error = "Processes {} exited with error code {}".format(
" ".join([str(i) for (i, _) in errored_processes]),
MultiProcessTestCase.TEST_ERROR_EXIT_CODE,
)
raise RuntimeError(error)
# If no process exited uncleanly, we check for timeouts, and then ensure
# each process exited cleanly.
for i, p in enumerate(self.processes):
if p.exitcode is None:
raise RuntimeError('Process {} terminated or timed out after {} seconds'.format(i, elapsed_time))
self.assertEqual(
p.exitcode,
first_process.exitcode,
msg="Expect process {} exit code to match Process 0 exit code of {}, but got {}".format(
i, first_process.exitcode, p.exitcode
),
)
for skip in TEST_SKIPS.values():
if first_process.exitcode == skip.exit_code:
raise unittest.SkipTest(skip.message)
self.assertEqual(
first_process.exitcode,
0,
msg="Expected zero exit code but got {}".format(first_process.exitcode)
)
@property
def is_master(self):
return self.rank == 0