done adding main logic in run_train_nerf to initialize processes with more robust handling of DDP

mattport/nerf/trainer.py

@@ -8,8 +8,6 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 from omegaconf import DictConfig
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.utils.data.distributed import DistributedSampler
 
 from mattport.nerf.dataset.build import build_dataset
 from mattport.nerf.optimizer import Optimizer
@@ -54,52 +52,16 @@ class Trainer:
         """
         raise NotImplementedError
 
-    def train(self, gpu: int) -> None:
-        """_summary_"""
-        self.setup_ddp(gpu, self.config.world_size)
-
-        # TODO(): implement training pipeline. below I add skeleton code for compatibility with DDP
-        # self.setup_graph()
-        # self.graph = DDP(self.graph, device_ids=[gpu])
-        # self.setup_dataset()
-        # train_sampler = DistributedSampler(self.dataset, rank=gpu, num_replicas=self.config.world_size, shuffle=True)
-        # TODO(): create dataloader that takes as input the train_sampler
-        # self.setup_optimizer()
-
-        # TODO(): calls to train_epoch and train_iteration loop
-
-        dist.destroy_process_group()
-        raise NotImplementedError
-
-    def train_epoch(self, gpu: int):
+    def train(self) -> None:
         """_summary_"""
         raise NotImplementedError
 
-    def train_iteration(self, gpu: int):
+    def train_epoch(self):
+        """_summary_"""
+        raise NotImplementedError
+
+    def train_iteration(self):
         """_summary_"""
         # TODO(): save checkpoints, and do logging only on rank 0 device only
         # if gpu == 0:
         raise NotImplementedError
-
-    def setup_ddp(self, gpu: int, world_size: int) -> None:
-        """Initialization of processes for distributed training.
-
-        Args:
-            gpu (int): process rank number
-            world_size (int): total number of gpus available for training
-        """
-        os.environ["MASTER_ADDR"] = "localhost"
-        os.environ["MASTER_PORT"] = "12356"
-        dist.init_process_group("nccl", init_method="env://", rank=gpu, world_size=world_size)
-        torch.cuda.set_device(gpu)
-
-        torch.manual_seed(0)
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
-        np.random.seed(0)
-
-    def run_multiprocess_train(self) -> None:
-        """function to spawn multiple processes for training."""
-        if self.config.world_size < 1:
-            self.config.world_size = torch.cuda.device_count()
-        mp.spawn(self.train, nprocs=self.config.world_size)

scripts/run_train_nerf.py

@@ -1,21 +1,190 @@
 """
 run_train_nerf.py
 """
+import datetime
 import logging
+import random
+import socket
+from datetime import timedelta
+from typing import Any, Callable
 
 import hydra
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
 from omegaconf import DictConfig
 
 from mattport.nerf.trainer import Trainer
 
 logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG)
 
+DEFAULT_TIMEOUT = timedelta(minutes=30)
+
+
+def _find_free_port() -> str:
+    """Find free port on open socket"""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(("", 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+def _set_random_seed(seed) -> None:
+    """Set randomness seed in torch and numpy"""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def _synchronize(world_size, local_rank):
+    """Wait for all processes on node finish"""
+    if world_size > 1:
+        dist.barrier(device_ids=[local_rank])
+
+
+def _distributed_worker(
+    local_rank: int,
+    main_func: Callable,
+    world_size: int,
+    num_gpus_per_machine: int,
+    machine_rank: int,
+    dist_url: str,
+    config: DictConfig,
+    timeout: datetime = DEFAULT_TIMEOUT,
+) -> Any:
+    """Spawned distributed worker that handles the initialization of process group and handles the
+       training process on multiple processes
+
+    Args:
+        local_rank (int): current rank of process
+        main_func (Callable): function that will be called by the distributed workers
+        world_size (int): total number of gpus available
+        num_gpus_per_machine (int): number of GPUs per machine
+        machine_rank (int): rank of this machine
+        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
+        config (DictConfig): config file specifying trainng regimen
+        timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT.
+
+    Raises:
+        e: Exception in initializing the process group
+
+    Returns:
+        Any: TODO(): determine the return type
+    """
+    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,
+            timeout=timeout,
+        )
+    except Exception as e:
+        logger = logging.getLogger(__name__)
+        logger.error("Process group URL: %s", dist_url)
+        raise e
+
+    assert num_gpus_per_machine <= torch.cuda.device_count()
+    torch.cuda.set_device(local_rank)
+    _set_random_seed(1234 + local_rank)
+    _synchronize(world_size, local_rank)
+
+    output = main_func(local_rank, world_size, config)
+    _synchronize(world_size, local_rank)
+    dist.destroy_process_group()
+    return output
+
+
+def _train(local_rank: int, world_size: int, config: DictConfig) -> Any:
+    """Main training function that sets up and runs the trainer per process
+
+    Args:
+        local_rank (int): _description_
+        world_size (int): _description_
+        config (DictConfig): _description_
+
+    Returns:
+        Any: TODO(): determine the return type
+    """
+    # trainer = Trainer(local_rank, world_size, config)
+    # trainer.train()
+    return 0
+
+
+def launch(
+    main_func: Callable,
+    num_gpus_per_machine: int,
+    num_machines: int = 1,
+    machine_rank: int = 0,
+    dist_url: str = "auto",
+    config: DictConfig = None,
+    timeout: datetime = DEFAULT_TIMEOUT,
+) -> None:
+    """Function that spawns muliple processes to call on main_func
+
+    Args:
+        main_func (Callable): function that will be called by the distributed workers
+        num_gpus_per_machine (int): number of GPUs per machine
+        num_machines (int, optional): total number of machines
+        machine_rank (int, optional): rank of this machine. Defaults to 0.
+        dist_url (str, optional): url to connect to for distributed jobs. Defaults to "auto".
+        config (DictConfig, optional): config file specifying trainng regimen Defaults to None.
+        timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT.
+    """
+    world_size = num_machines * num_gpus_per_machine
+    if world_size > 1:
+        if dist_url == "auto":
+            assert num_machines == 1, "dist_url=auto not supported in multi-machine jobs."
+            port = _find_free_port()
+            dist_url = f"tcp://127.0.0.1:{port}"
+        if num_machines > 1 and dist_url.startswith("file://"):
+            logger = logging.getLogger(__name__)
+            logger.warning("file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://")
+
+        process_context = mp.spawn(
+            _distributed_worker,
+            nprocs=num_gpus_per_machine,
+            args=(
+                main_func,
+                world_size,
+                num_gpus_per_machine,
+                machine_rank,
+                dist_url,
+                config,
+                timeout,
+            ),
+            join=False,
+        )
+        try:
+            process_context.join()
+        except KeyboardInterrupt:
+            for i, process in enumerate(process_context.processes):
+                if process.is_alive():
+                    logger.info("Terminating process %s", str(i))
+                    process.terminate()
+                process.join()
+                logger.info("Process %s finished", str(i))
+    else:
+        main_func(config)
+
 
 @hydra.main(config_path="../configs", config_name="config.yaml")
 def main(config: DictConfig):
     """Main function."""
-    trainer = Trainer(config)
-    trainer.run_multiprocess_train()
+    launch(
+        _train,
+        config.machine_config.num_gpus,
+        num_machines=config.machine_config.num_machines,
+        machine_rank=config.machine_config.machine_rank,
+        dist_url=config.machine_config.dist_url,
+        config=(config,),
+    )
 
 
 if __name__ == "__main__":

tests/nerf/configs/test_basic_graph.yaml

@@ -1,4 +1,8 @@
-world_size: 0
+machine_config:
+  num_gpus: 8
+  num_machines: 1
+  machine_rank: 0
+  dist_url: auto
 network:
   mlp_0:
     class_name: MLP