21 lines
650 B
Python
21 lines
650 B
Python
import time
|
|
|
|
import torch
|
|
|
|
|
|
def custom_barrier_with_timeout(timeout_sec=60 * 60 * 2, check_interval=60 * 5):
|
|
"""
|
|
Repeatedly attempts to synchronize processes using torch.distributed.barrier().
|
|
Retries until timeout_sec is exceeded.
|
|
"""
|
|
start_time = time.time()
|
|
while True:
|
|
try:
|
|
torch.distributed.barrier()
|
|
break # Success
|
|
except Exception as e:
|
|
elapsed = time.time() - start_time
|
|
if elapsed > timeout_sec:
|
|
raise TimeoutError(f"Barrier timed out after {timeout_sec} seconds") from e
|
|
time.sleep(check_interval) # Wait before retrying
|