| 510 | |
| 511 | @classmethod |
| 512 | def sync(cls, wait_for=None, timeout=10): |
| 513 | if wait_for is None: |
| 514 | wait_for = dist.get_world_size() |
| 515 | cls.barrier_id += 1 |
| 516 | barrier_dir = os.path.join(os.environ["TEMP_DIR"], "barrier") |
| 517 | pid = str(os.getpid()) |
| 518 | barrier_file = os.path.join(barrier_dir, pid) |
| 519 | with _lock(): |
| 520 | with open(barrier_file, "w") as f: |
| 521 | f.write(str(cls.barrier_id)) |
| 522 | |
| 523 | start_time = time.time() |
| 524 | while True: |
| 525 | arrived = 0 |
| 526 | with _lock(): |
| 527 | for f_name in os.listdir(barrier_dir): |
| 528 | with open(os.path.join(barrier_dir, f_name)) as f: |
| 529 | data = f.read() |
| 530 | if int(data) >= cls.barrier_id: |
| 531 | arrived += 1 |
| 532 | if arrived == wait_for: |
| 533 | break |
| 534 | |
| 535 | if time.time() - start_time > timeout: |
| 536 | raise RuntimeError("barrier timeout") |
| 537 | time.sleep(0.1) |
| 538 | |
| 539 | |
| 540 | class TestDistBackend(MultiProcessTestCase): |