Pass NCCL timeout as an argument to launch
Summary: Pass NCCL timeout as an argument to launch
# Copyright (c) Facebook, Inc. and its affiliates.
import logging
+from datetime import timedelta
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from detectron2.utils import comm
-__all__ = ["launch"]
+__all__ = ["DEFAULT_TIMEOUT", "launch"]
+
+DEFAULT_TIMEOUT = timedelta(minutes=30)
def _find_free_port():
@@ -21,7 +24,15 @@ def _find_free_port():
return port
-def launch(main_func, num_gpus_per_machine, num_machines=1, machine_rank=0, dist_url=None, args=()):
+def launch(
+ main_func,
+ num_gpus_per_machine,
+ num_machines=1,
+ machine_rank=0,
+ dist_url=None,
+ args=(),
+ timeout=DEFAULT_TIMEOUT,
+):
"""
Launch multi-gpu or distributed training.
This function must be called on all machines involved in the training.
@@ -35,6 +46,7 @@ def launch(main_func, num_gpus_per_machine, num_machines=1, machine_rank=0, dist
dist_url (str): url to connect to for distributed jobs, including protocol
e.g. "tcp://127.0.0.1:8686".
Can be set to "auto" to automatically select a free port on localhost
+ timeout (timedelta): timeout of the distributed workers
args (tuple): arguments passed to main_func
"""
world_size = num_machines * num_gpus_per_machine
@@ -55,7 +67,15 @@ def launch(main_func, num_gpus_per_machine, num_machines=1, machine_rank=0, dist
mp.spawn(
_distributed_worker,
nprocs=num_gpus_per_machine,
- args=(main_func, world_size, num_gpus_per_machine, machine_rank, dist_url, args),
+ args=(
+ main_func,
+ world_size,
+ num_gpus_per_machine,
+ machine_rank,
+ dist_url,
+ args,
+ timeout,
+ ),
daemon=False,
)
else:
@@ -63,13 +83,24 @@ def launch(main_func, num_gpus_per_machine, num_machines=1, machine_rank=0, dist
def _distributed_worker(
- local_rank, main_func, world_size, num_gpus_per_machine, machine_rank, dist_url, args
+ local_rank,
+ main_func,
+ world_size,
+ num_gpus_per_machine,
+ machine_rank,
+ dist_url,
+ args,
+ timeout=DEFAULT_TIMEOUT,
):
assert torch.cuda.is_available(), "cuda is not available. Please check your installation."
global_rank = machine_rank * num_gpus_per_machine + local_rank
try:
dist.init_process_group(
- backend="NCCL", init_method=dist_url, world_size=world_size, rank=global_rank
+ backend="NCCL",
+ init_method=dist_url,
+ world_size=world_size,
+ rank=global_rank,
+ timeout=timeout,
)
except Exception as e:
logger = logging.getLogger(__name__)
