()
| 22 | |
| 23 | |
| 24 | def main(): |
| 25 | argslist = list(sys.argv)[1:] |
| 26 | world_size = torch.cuda.device_count() |
| 27 | |
| 28 | if '--world-size' in argslist: |
| 29 | argslist[argslist.index('--world-size') + 1] = str(world_size) |
| 30 | else: |
| 31 | argslist.append('--world-size') |
| 32 | argslist.append(str(world_size)) |
| 33 | |
| 34 | workers = [] |
| 35 | |
| 36 | for i in range(world_size): |
| 37 | if '--rank' in argslist: |
| 38 | argslist[argslist.index('--rank') + 1] = str(i) |
| 39 | else: |
| 40 | argslist.append('--rank') |
| 41 | argslist.append(str(i)) |
| 42 | stdout = None if i == 0 else subprocess.DEVNULL |
| 43 | worker = subprocess.Popen( |
| 44 | [str(sys.executable)] + argslist, stdout=stdout) |
| 45 | workers.append(worker) |
| 46 | |
| 47 | returncode = 0 |
| 48 | try: |
| 49 | pending = len(workers) |
| 50 | while pending > 0: |
| 51 | for worker in workers: |
| 52 | try: |
| 53 | worker_returncode = worker.wait(1) |
| 54 | except subprocess.TimeoutExpired: |
| 55 | continue |
| 56 | pending -= 1 |
| 57 | if worker_returncode != 0: |
| 58 | if returncode != 1: |
| 59 | for worker in workers: |
| 60 | worker.terminate() |
| 61 | returncode = 1 |
| 62 | |
| 63 | except KeyboardInterrupt: |
| 64 | print('Pressed CTRL-C, TERMINATING') |
| 65 | for worker in workers: |
| 66 | worker.terminate() |
| 67 | for worker in workers: |
| 68 | worker.wait() |
| 69 | raise |
| 70 | |
| 71 | sys.exit(returncode) |
| 72 | |
| 73 | |
| 74 | if __name__ == "__main__": |
no test coverage detected